LeetCode – Number of Squareful Arrays (Java)

by

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

Java Solution

This problem can be solved by using the similar method of Permutation II. A set can be used to track if same element is swapped. The major difference is that we need to check if previous adjacent two neighbors are squareful. 

class Solution {
    int count = 0;
 
    public int numSquarefulPerms(int[] A) {
        Arrays.sort(A);
        helper(A, 0);   
        return count;    
    }
 
    void helper(int[] A, int start){
        //the adjacent two numbers before start
        if(start>1 && (!isSquareful(A[start], A[start-1]) || !isSquareful(A[start-1], A[start-2]))){
            return;
        }
        //if start is the last, then check start with its adjancent neighbor
        if(start==A.length-1 && !isSquareful(A[start], A[start-1])){
            return;
        }
 
        if(start==A.length-1){
            count++;
            return;
        }
 
        HashSet<Integer> set = new HashSet<>();
        for(int i=start; i<A.length; i++){
            //use set to track if the same element is used
            if(set.contains(A[i])){
                continue;
            }
            set.add(A[i]);
 
            swap(A, i, start);
            helper(A, start+1);
            swap(A, i, start);
        }
    }
 
    void swap(int[] A, int i, int j){
        int t = A[i];
        A[i] = A[j];
        A[j] = t;
    }
 
    boolean isSquareful(int a, int b){
        double root = Math.sqrt(a+b);
        return (root-Math.floor(root))==0;
    }
}

Leave a Comment