Medium
Count the Number of Square-Free Subsets — C++
Full explanation · Time O(n + m * 2^p) · Space O(m * 2^p)
// Time: O(n + m * 2^p)
// Space: O(m * 2^p)
// number theory, combinatorics, bitmasks, dp
class Solution {
public:
int squareFreeSubsets(vector<int>& nums) {
const int MAX_NUM = *max_element(cbegin(nums), cend(nums));
const auto& linear_sieve_of_eratosthenes = [](int n) { // Time: O(n), Space: O(n)
vector<int> spf(n + 1, -1);
vector<int> primes;
for (int i = 2; i <= n; ++i) {
if (spf[i] == -1) {
spf[i] = i;
primes.emplace_back(i);
}
for (const auto& p : primes) {
if (i * p > n || p > spf[i]) {
break;
}
spf[i * p] = p;
}
}
return primes;
};
const auto& PRIMES = linear_sieve_of_eratosthenes(MAX_NUM);
vector<int> MASKS(MAX_NUM + 1);
for (int x = 0; x <= MAX_NUM; ++x) {
for (int i = 0, y = x; i < size(PRIMES); ++i) {
if (y % PRIMES[i]) {
continue;
}
if (y % (PRIMES[i] * PRIMES[i]) == 0) {
MASKS[x] = 0;
break;
}
MASKS[x] |= (1 << i);
y /= PRIMES[i];
}
}
static const int MOD = 1e9 + 7;
const auto& powmod = [&](int a, int b) {
a %= MOD;
int64_t result = 1;
while (b) {
if (b & 1) {
result = result * a % MOD;
}
a = int64_t(a) * a % MOD;
b >>= 1;
}
return result;
};
unordered_map<int, int> cnt;
for (const auto& x : nums) {
++cnt[x];
}
vector<int> arr;
for (const auto& [k, _] : cnt) {
if (k != 1) {
arr.emplace_back(k);
}
}
vector<int> dp(1 << size(PRIMES), 1);
for (const auto& x : arr) {
if (!MASKS[x]) {
continue;
}
for (int mask = size(dp) - 1; mask >= 0; --mask) {
if ((MASKS[x] & mask) == 0) {
dp[mask | MASKS[x]] = (dp[mask | MASKS[x]] + static_cast<int64_t>(cnt[x]) * dp[mask]) % MOD;
}
}
}
return cnt.count(1) ? ((dp.back() * powmod(2, cnt[1]) - 1) + MOD) % MOD: ((dp.back() - 1) + MOD) % MOD;
}
};
// Time: O(n + m * 2^p)
// Space: O(m * 2^p)
// number theory, combinatorics, bitmasks, memoization
class Solution2 {
public:
int squareFreeSubsets(vector<int>& nums) {
const int MAX_NUM = *max_element(cbegin(nums), cend(nums));
const auto& linear_sieve_of_eratosthenes = [](int n) { // Time: O(n), Space: O(n)
vector<int> spf(n + 1, -1);
vector<int> primes;
for (int i = 2; i <= n; ++i) {
if (spf[i] == -1) {
spf[i] = i;
primes.emplace_back(i);
}
for (const auto& p : primes) {
if (i * p > n || p > spf[i]) {
break;
}
spf[i * p] = p;
}
}
return primes;
};
const auto& PRIMES = linear_sieve_of_eratosthenes(MAX_NUM);
vector<int> MASKS(MAX_NUM + 1);
for (int x = 0; x <= MAX_NUM; ++x) {
for (int i = 0, y = x; i < size(PRIMES); ++i) {
if (y % PRIMES[i]) {
continue;
}
if (y % (PRIMES[i] * PRIMES[i]) == 0) {
MASKS[x] = 0;
break;
}
MASKS[x] |= (1 << i);
y /= PRIMES[i];
}
}
static const int MOD = 1e9 + 7;
const auto& powmod = [&](int a, int b) {
a %= MOD;
int64_t result = 1;
while (b) {
if (b & 1) {
result = result * a % MOD;
}
a = int64_t(a) * a % MOD;
b >>= 1;
}
return result;
};
unordered_map<int, int> cnt;
for (const auto& x : nums) {
++cnt[x];
}
vector<int> arr;
for (const auto& [k, _] : cnt) {
if (k != 1) {
arr.emplace_back(k);
}
}
vector<vector<int>> dp(size(arr), vector<int>(1 << size(PRIMES), -1));
const function<int(int, int)> memoization = [&](int i, int mask) {
if (i == size(arr)) {
return 1;
}
if (dp[i][mask] == -1) {
dp[i][mask] = memoization(i + 1, mask);
if (MASKS[arr[i]] && (MASKS[arr[i]] & mask) == MASKS[arr[i]]) {
dp[i][mask] = (dp[i][mask] + static_cast<int64_t>(cnt[arr[i]]) * memoization(i + 1, mask ^ MASKS[arr[i]])) % MOD;
}
}
return dp[i][mask];
};
return cnt.count(1) ? ((memoization(0, (1 << size(PRIMES)) - 1) * powmod(2, cnt[1]) - 1) + MOD) % MOD: ((memoization(0, (1 << size(PRIMES)) - 1) - 1) + MOD) % MOD;
}
};