Leetcode 996. Number of Squareful Arrays

This problem is asking to return the number of permutations of an array that are "squareful". An array (or permutation of the array) is defined as "squareful" if the sum of every pair of adjacent elements is a perfect square.

For example, let's consider the array [1, 17, 8]. There are three distinct permutations of the array which are:

  1. [1, 8, 17]
  2. [1, 17, 8]
  3. [8, 1, 17]
  4. [8, 17, 1]
  5. [17, 1, 8]
  6. [17, 8, 1]

In the above permutations, only [1, 8, 17] and [17, 8, 1] are "squareful" because their adjacent pair sums (1+8=9, 8+17=25 and 17+8=25, 8+1=9) are both perfect squares. So, the output should be 2.

The solution is using a depth-first search (DFS) to traverse through all the possible permutations of the array and it uses the "used" array to keep track of the elements that have been used in the current permutation. Before adding an element to the permutation, it checks if the sum of this element and the last element in the current permutation is a perfect square. If not, it skips this element and continues to the next one. If it is a perfect square, it adds this element to the permutation and recursively calls the dfs function. After the dfs function returns, it removes this element from the permutation and marks it as unused.

It also optimizes the solution by skipping duplicate elements. If the current element is the same as the previous one and the previous one is not used, it continues to the next one.

Finally, when the size of the permutation is equal to the size of the original array, it means that a valid permutation is found, so it increments the answer by one.

Here is the solution implemented in Python:

1
2python
3import math
4class Solution:
5    def numSquarefulPerms(self, A) -> int:
6        counter, res, N = collections.Counter(A), 0, len(A)
7
8        graph = {x: [] for x in counter}
9        for x in counter:
10            for y in counter:
11                if int((x+y)**.5 + 0.5) ** 2 == x+y:
12                    graph[x].append(y)
13
14        def dfs(x, todo):
15            counter[x] -= 1
16            if todo == 0:
17                ans = 1
18            else:
19                ans = sum(dfs(y, todo - 1) for y in graph[x] if counter[y])
20            counter[x] += 1
21            return ans
22
23        return sum(dfs(x, len(A) - 1) for x in counter)

It uses a counter to count the occurrence of each number, a graph to keep track of which pairs can make a square, and a DFS to traverse the graph.Here is the solution implemented in JavaScript:

1
2javascript
3var numSquarefulPerms = function(A) {
4    A.sort((a, b) => a - b);
5    let res = [0];
6    dfs(A, [], Array(A.length).fill(false), res);
7    return res[0];
8};
9
10var dfs = function(A, path, used, res) {
11    if (path.length === A.length) {
12        res[0]++;
13        return;
14    }
15    for (let i = 0; i < A.length; i++) {
16        if (used[i] || (i > 0 && A[i] === A[i - 1] && !used[i - 1])) continue;
17        if (path.length > 0 && !isSquare(A[i] + path[path.length - 1])) continue;
18        used[i] = true;
19        path.push(A[i]);
20        dfs(A, path, used, res);
21        path.pop();
22        used[i] = false;
23    }
24};
25
26var isSquare = function(n) {
27    let sqrt = Math.sqrt(n);
28    return sqrt === Math.floor(sqrt);
29};

In the JavaScript code, isSquare function checks whether a number is a perfect square by taking the square root of a number and checking if it is an integer. DFS, which stands for depth-first search, checks all possible permutations of array and increments the result if the permutation is squareful.

Here is the solution implemented in Java:

1
2java
3class Solution {
4    public int numSquarefulPerms(int[] A) {
5        Arrays.sort(A);
6        return dfs(A, new int[A.length], new boolean[A.length], 0, -1);
7    }
8
9    public int dfs(int[] A, int[] path, boolean[] used, int idx, int prev) {
10        if (idx == A.length) return 1;
11        int res = 0;
12        for (int i = 0; i < A.length; i++) {
13            if (used[i]) continue;
14            if (i > 0 && A[i] == A[i - 1] && !used[i - 1]) continue;           
15            if (prev != -1 && !isSquare(A[i] + prev)) continue;           
16            path[idx] = A[i];
17            used[i] = true;
18            res += dfs(A, path, used, idx + 1, A[i]);
19            used[i] = false;
20        }
21        return res;
22    }
23
24    public boolean isSquare(int n) {
25        int sqr = (int) Math.sqrt(n);
26        return sqr * sqr == n;
27    }
28}

In the Java code, isSquare function is a helper function that checks whether a number is a perfect square. The function numSquarefulPerms applies the depth-first search (DFS) by calling dfs function and handles the different cases of array permutations.


Got a question? Ask the Teaching Assistant anything you don't understand.

Still not clear? Ask in the Forum,  Discord or Submit the part you don't understand to our editors.


TA 👨‍🏫