Skip to content

2597. The Number of Beautiful Subsets 👍

  • Time: $O(n\log n + k)$
  • Space: $O(n + k)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
// e.g. nums = [2, 3, 4, 4], k = 2
//
// subset[0] = [2, 4, 4']
// subset[1] = [1]
// count = {2: 1, 4: 2, 1: 1}
//
// Initially, skip = len([]) = 0, pick = len([]) = 0
//
// * For values in subset[0]:
//   After 2:
//     skip = skip + pick = len([]) = 0
//     pick = (2^count[2] - 1) * (1 + skip + pick)
//          = len([[2]]) * len([[]])
//          = len([[2]]) = 1
//   After 4:
//     skip = skip + pick = len([[2]]) = 1
//     pick = (2^count[4] - 1) * (1 + skip)
//          = len([[4], [4'], [4, 4']]) * len([[]])
//          = len([[4], [4'], [4, 4']]) = 3
//
// * For values in subset[1]:
//   After 1:
//     skip = skip + pick
//          = len([[2], [4], [4'], [4, 4']]) = 4
//     pick = (2^count[1] - 1) * (1 + skip + pick)
//          = len([[1]]) * len([[], [2], [4], [4'], [4, 4']])
//          = len([[1], [1, 2], [1, 4], [1, 4'], [1, 4, 4']]) = 5
//
// So, ans = skip + pick = 9

class Solution {
 public:
  int beautifulSubsets(vector<int>& nums, int k) {
    constexpr int kMaxNum = 1000;
    vector<int> count(kMaxNum + 1);
    unordered_map<int, set<int>> modToSubset;

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

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

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

    return 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
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
// e.g. nums = [2, 3, 4, 4], k = 2
//
// subset[0] = [2, 4, 4']
// subset[1] = [1]
// count = {2: 1, 4: 2, 1: 1}
//
// Initially, skip = len([]) = 0, pick = len([]) = 0
//
// * For values in subset[0]:
//   After 2:
//     skip = skip + pick = len([]) = 0
//     pick = (2^count[2] - 1) * (1 + skip + pick)
//          = len([[2]]) * len([[]])
//          = len([[2]]) = 1
//   After 4:
//     skip = skip + pick = len([[2]]) = 1
//     pick = (2^count[4] - 1) * (1 + skip)
//          = len([[4], [4'], [4, 4']]) * len([[]])
//          = len([[4], [4'], [4, 4']]) = 3
//
// * For values in subset[1]:
//   After 1:
//     skip = skip + pick
//          = len([[2], [4], [4'], [4, 4']]) = 4
//     pick = (2^count[1] - 1) * (1 + skip + pick)
//          = len([[1]]) * len([[], [2], [4], [4'], [4, 4']])
//          = len([[1], [1, 2], [1, 4], [1, 4'], [1, 4, 4']]) = 5
//
// So, ans = skip + pick = 9

class Solution {
  public int beautifulSubsets(int[] nums, int k) {
    final int kMaxNum = 1000;
    int[] count = new int[kMaxNum + 1];
    Map<Integer, Set<Integer>> modToSubset = new HashMap<>();

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

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

    for (Set<Integer> subset : modToSubset.values())
      for (final int num : subset) {
        final int nonEmptyCount = (int) Math.pow(2, count[num]) - 1;
        final int cacheSkip = skip;
        skip += pick;
        pick = nonEmptyCount * (1 + cacheSkip + (num - prevNum == k ? 0 : pick));
        prevNum = num;
      }

    return 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
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
# e.g. nums = [2, 3, 4, 4], k = 2
#
# subset[0] = [2, 4, 4']
# subset[1] = [1]
# count = {2: 1, 4: 2, 1: 1}
#
# Initially, skip = len([]) = 0, pick = len([]) = 0
#
# * For values in subset[0]:
#   After 2:
#     skip = skip + pick = len([]) = 0
#     pick = (2^count[2] - 1) * (1 + skip + pick)
#          = len([[2]]) * len([[]])
#          = len([[2]]) = 1
#   After 4:
#     skip = skip + pick = len([[2]]) = 1
#     pick = (2^count[4] - 1) * (1 + skip)
#          = len([[4], [4'], [4, 4']]) * len([[]])
#          = len([[4], [4'], [4, 4']]) = 3
#
# * For values in subset[1]:
#   After 1:
#     skip = skip + pick
#          = len([[2], [4], [4'], [4, 4']]) = 4
#     pick = (2^count[1] - 1) * (1 + skip + pick)
#          = len([[1]]) * len([[], [2], [4], [4'], [4, 4']])
#          = len([[1], [1, 2], [1, 4], [1, 4'], [1, 4, 4']]) = 5
#
# So, ans = skip + pick = 9

class Solution:
  def beautifulSubsets(self, nums: list[int], k: int) -> int:
    count = collections.Counter(nums)
    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):
        nonEmptyCount = 2**count[num] - 1
        skip, pick = (skip + pick,
                      nonEmptyCount *
                      (1 + skip + (0 if num - prevNum == k else pick)))
        prevNum = num

    return skip + pick