Skip to content

1994. The Number of Good Subsets 👍

  • Time: $O(\max(\texttt{nums}) \cdot 2^{10})$
  • Space: $O(\max(\texttt{nums}) + 2^{10})$
 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
class Solution {
 public:
  int numberOfGoodSubsets(vector<int>& nums) {
    const vector<int> primes{2, 3, 5, 7, 11, 13, 17, 19, 23, 29};
    const int n = 1 << primes.size();
    const int maxNum = ranges::max(nums);
    vector<long> dp(n);
    vector<int> count(maxNum + 1);

    dp[0] = 1;

    for (const int num : nums)
      ++count[num];

    for (int num = 2; num <= maxNum; ++num) {
      if (count[num] == 0)
        continue;
      if (num % 4 == 0 || num % 9 == 0 || num % 25 == 0)
        continue;
      const int numPrimesMask = getPrimesMask(num, primes);
      for (int primesMask = 0; primesMask < n; ++primesMask) {
        if ((primesMask & numPrimesMask) > 0)
          continue;
        const int nextPrimesMask = primesMask | numPrimesMask;
        dp[nextPrimesMask] += dp[primesMask] * count[num];
        dp[nextPrimesMask] %= kMod;
      }
    }

    return modPow(2, count[1]) *
           (accumulate(dp.begin() + 1, dp.end(), 0L) % kMod) % kMod;
  }

 private:
  static constexpr int kMod = 1'000'000'007;

  int getPrimesMask(int num, const vector<int>& primes) {
    int primesMask = 0;
    for (int i = 0; i < primes.size(); ++i)
      if (num % primes[i] == 0)
        primesMask |= 1 << i;
    return primesMask;
  }

  long modPow(long x, long n) {
    if (n == 0)
      return 1;
    if (n % 2 == 1)
      return x * modPow(x % kMod, (n - 1)) % kMod;
    return modPow(x * x % kMod, (n / 2)) % kMod;
  }
};

 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
class Solution {
  public int numberOfGoodSubsets(int[] nums) {
    final int[] primes = {2, 3, 5, 7, 11, 13, 17, 19, 23, 29};
    final int n = 1 << primes.length;
    final int maxNum = Arrays.stream(nums).max().getAsInt();
    long[] dp = new long[n];
    int[] count = new int[maxNum + 1];

    dp[0] = 1;

    for (final int num : nums)
      ++count[num];

    for (int num = 2; num <= maxNum; ++num) {
      if (count[num] == 0)
        continue;
      if (num % 4 == 0 || num % 9 == 0 || num % 25 == 0)
        continue;
      final int numPrimesMask = getPrimesMask(num, primes);
      for (int primesMask = 0; primesMask < n; ++primesMask) {
        if ((primesMask & numPrimesMask) > 0)
          continue;
        final int nextPrimesMask = primesMask | numPrimesMask;
        dp[nextPrimesMask] += dp[primesMask] * count[num];
        dp[nextPrimesMask] %= kMod;
      }
    }

    return (int) (modPow(2, count[1]) * ((Arrays.stream(dp).sum() - 1) % kMod) % kMod);
  }

  private static final int kMod = 1_000_000_007;

  private int getPrimesMask(int num, int[] primes) {
    int primesMask = 0;
    for (int i = 0; i < primes.length; ++i)
      if (num % primes[i] == 0)
        primesMask |= 1 << i;
    return primesMask;
  }

  private long modPow(long x, long n) {
    if (n == 0)
      return 1;
    if (n % 2 == 1)
      return x * modPow(x, n - 1) % kMod;
    return modPow(x * x % kMod, n / 2);
  }
}

 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
class Solution:
  def numberOfGoodSubsets(self, nums: list[int]) -> int:
    kMod = 1_000_000_007
    primes = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29]
    n = 1 << len(primes)
    # dp[i] := the number of good subsets with set of primes = i bit mask
    dp = [1] + [0] * (n - 1)
    count = collections.Counter(nums)

    for num, freq in count.items():
      if num == 1:
        continue
      if any(num % squared == 0 for squared in [4, 9, 25]):
        continue
      numPrimesMask = sum(1 << i
                          for i, prime in enumerate(primes)
                          if num % prime == 0)
      for primesMask in range(n):
        # Skip since there're commen set of primes (becomes invalid subset)
        if primesMask & numPrimesMask > 0:
          continue
        nextPrimesMask = numPrimesMask | primesMask
        dp[nextPrimesMask] += dp[primesMask] * freq
        dp[nextPrimesMask] %= kMod

    return (1 << count[1]) * sum(dp[1:]) % kMod