Skip to content

2638. Count the Number of K-Free Subsets 👍

  • Time: $O(n\log n)$
  • Space: $O(n)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class Solution {
 public:
  // Similar to 2597. The Number of Beautiful Subsets
  long long countTheNumOfKFreeSubsets(vector<int>& nums, int k) {
    unordered_map<int, set<int>> modToSubset;

    for (const int num : nums)
      modToSubset[num % k].insert(num);

    int prevNum = -k;
    long skip = 0;
    long pick = 0;

    for (const auto& [_, subset] : modToSubset)
      for (const int num : subset) {
        const long cacheSkip = skip;
        skip += pick;
        pick = 1 + cacheSkip + (num - prevNum == k ? 0 : pick);
        prevNum = num;
      }

    return 1 + skip + pick;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class Solution {
  public long countTheNumOfKFreeSubsets(int[] nums, int k) {
    Map<Integer, Set<Integer>> modToSubset = new HashMap<>();

    for (final int num : nums) {
      modToSubset.putIfAbsent(num % k, new TreeSet<>());
      modToSubset.get(num % k).add(num);
    }

    int prevNum = -k;
    long skip = 0;
    long pick = 0;

    for (Set<Integer> subset : modToSubset.values())
      for (final int num : subset) {
        final long cacheSkip = skip;
        skip += pick;
        pick = 1 + cacheSkip + (num - prevNum == k ? 0 : pick);
        prevNum = num;
      }

    return 1 + skip + pick;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
class Solution:
  def countTheNumOfKFreeSubsets(self, nums: list[int], k: int) -> int:
    modToSubset = collections.defaultdict(set)

    for num in nums:
      modToSubset[num % k].add(num)

    prevNum = -k
    skip = 0
    pick = 0

    for subset in modToSubset.values():
      for num in sorted(subset):
        skip, pick = (skip + pick,
                      1 + skip + (0 if num - prevNum == k else pick))
        prevNum = num

    return 1 + skip + pick