Skip to content

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:

Input: [1,17,8]
Output: 2
Explanation:
[1,8,17] and [17,8,1] are the valid permutations.

Example 2:

Input: [2,2,2]
Output: 1

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;
  }
};