996

996. Number of Squareful Arrays (Hard)

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

 

Note:

1. 1 <= A.length <= 12
2. 0 <= A[i] <= 1e9

Related Topics:
Math, Backtracking, Graph

Similar Questions:

• Permutations II (Medium)

Solution 1. Backtracking

Brute force: next_permutation + squareful check.

Issue: if the leading x elements are not squareful, no matter how the rest of the elements are permutated, the result can't be squareful. So we need to make sure the permutation is squareful along the way.

Solution: manually implement next_permutation, check if the currently settled subarray is squareful. Skip this leading section if it's already not squareful.

// OJ: https://leetcode.com/problems/number-of-squareful-arrays/
// Author: github.com/lzl124631x
// Time: O(N!)
// Space: O(N^2)
class Solution {
    int ans = 0;
    inline bool isSquare(int n) {
        int r = sqrt(n);
        return r * r == n;
    }
    void dfs(vector<int> &A, int start) {
        if (start == A.size()) {
            ++ans;
            return;
        }
        unordered_set<int> s;
        for (int i = start; i < A.size(); ++i) {
            if (s.count(A[i]) || (start - 1 >= 0 && !isSquare(A[start - 1] + A[i]))) continue;
            s.insert(A[i]);
            swap(A[start], A[i]);
            dfs(A, start + 1);
            swap(A[start], A[i]);
        }
    }
public:
    int numSquarefulPerms(vector<int>& A) {
        dfs(A, 0);
        return ans;
    }
};