996 Number of Squareful Arrays¶
Given an array A of non-negative integers, the array is squareful if for every pair of adjacent elements, their sum is a perfect square.
Return the number of permutations of A that are squareful. Two permutations A1 and A2 differ if and only if there is some index i such that A1[i] != A2[i].
Example 1:
Example 2:
Solution 1 : DFS¶
Try all permutations with pruning.
Time complexity: O(n!) Space complexity: O(n)
//running time: 4 ms, 8.8 MB
class Solution {
public:
int numSquarefulPerms(vector<int>& A) {
std::sort(begin(A), end(A));
vector<int> cur;
vector<int> used(A.size());
int ans = 0;
dfs(A, cur, used, ans);
return ans;
}
private:
bool squareful(int x, int y) {
int s = sqrt(x + y);
return s * s == x + y;
}
void dfs(const vector<int>& A, vector<int>& cur, vector<int>& used, int& ans) {
if (cur.size() == A.size()) {
++ans;
return;
}
for (int i = 0; i < A.size(); ++i) {
if (used[i]) continue;
// Avoid duplications.
if (i > 0 && !used[i - 1] && A[i] == A[i - 1]) continue;
// Prune invalid solutions.
if (!cur.empty() && !squareful(cur.back(), A[i])) continue;
cur.push_back(A[i]);
used[i] = 1;
dfs(A, cur, used, ans);
used[i] = 0;
cur.pop_back();
}
}
};
Solution 2: DP Hamiltonian Path¶
dp[s][i] := # of ways to reach state s (binary mask of nodes visited) that ends with node i
dp[s | (1 << j)][j] += dp[s][i] if g[i][j]
Time complexity: \(O(n^2*2^n)\) Space complexity: \(O(2^n)\)
// running time: 20 ms, 14.9 MB
class Solution {
public:
int numSquarefulPerms(vector<int>& A) {
const int n = A.size();
// For deduplication.
std::sort(begin(A), end(A));
// g[i][j] == 1 if A[i], A[j] are squareful.
vector<vector<int>> g(n, vector<int>(n));
// dp[s][i] := number of ways to reach state s and ends with node i.
vector<vector<int>> dp(1 << n, vector<int>(n));
for (int i = 0; i < n; ++i) {
for (int j = 0; j < n; ++j) {
if (i == j) continue;
int r = sqrt(A[i] + A[j]);
if (r * r == A[i] + A[j])
g[i][j] = 1;
}
}
// For the same numbers, only the first one can be the starting point.
for (int i = 0; i < n; ++i)
if (i == 0 || A[i] != A[i - 1])
dp[(1 << i)][i] = 1;
int ans = 0;
for (int s = 0; s < (1 << n); ++s)
for (int i = 0; i < n; ++i) {
if (!dp[s][i]) continue;
for (int j = 0; j < n; ++j) {
if (!g[i][j]) continue;
if (s & (1 << j)) continue;
// Only the first one can be used as the dest.
if (j > 0 && !(s & (1 << (j - 1))) && A[j - 1] == A[j]) continue;
dp[s | (1 << j)][j] += dp[s][i];
}
}
for (int i = 0; i < n; ++i)
ans += dp[(1 << n) - 1][i];
return ans;
}
};