996. Number of Squareful Arrays

Question:

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 <= A.length <= 12
• 0 <= A[i] <= 1e9

Solution:

• Method 1: search O(N!) AC 3ms
  • Pruning is super important, I cannot AC without pruning.
  • need to check if the last two variables are squareful.

      class Solution {
      private int res = 0;
      public int numSquarefulPerms(int[] A) {
          if(A.length == 1) return 0;
          Arrays.sort(A);
          dfs(A, new ArrayList<Integer>(), 0);
          return res;
      }
      private void dfs(int[] A, List<Integer> temp, int visited){
          if(temp.size() == A.length){
              boolean flag = true;
              for(int i = 1; i < A.length; i++){               
                  flag &= square(temp.get(i - 1) + temp.get(i));
              }
              if(flag) res++;
          }else if(temp.size() <= 1 || square(temp.get(temp.size() - 1) + temp.get(temp.size() - 2))){
              for(int i = 0; i < A.length; i++){
                  if((visited & (1 << i)) > 0) continue;
                  if(i > 0 && A[i] == A[i - 1] && (visited & (1 << (i - 1))) == 0) continue;
                  temp.add(A[i]);
                  dfs(A, temp, visited | (1 << i));
                  temp.remove(temp.size() - 1);
              }
          }
      }
      private boolean square(int sum){
          return (int)Math.sqrt(sum) * (int)Math.sqrt(sum) == sum;
      }
      }
    

• Method 2: DP AC 99.23%
  • dp[i][state]: i means current index and state is a bit map show what are the next index to visit(which are set to 1).
  • dp[i][state] = dp[next][next bit is set to zero state] for all non-visited indice.
  • Result: sum of all index as head, all other index are not visited state.

      class Solution {
      private int res = 0;
      public int numSquarefulPerms(int[] A) {
          if(A.length == 1) return 0;
          Arrays.sort(A);
          dfs(A, new ArrayList<Integer>(), 0);
          return res;
      }
      private void dfs(int[] A, List<Integer> temp, int visited){
          if(temp.size() == A.length){
              boolean flag = true;
              for(int i = 1; i < A.length; i++){               
                  flag &= square(temp.get(i - 1) + temp.get(i));
              }
              if(flag) res++;
          }else if(temp.size() <= 1 || square(temp.get(temp.size() - 1) + temp.get(temp.size() - 2))){
              for(int i = 0; i < A.length; i++){
                  if((visited & (1 << i)) > 0) continue;
                  if(i > 0 && A[i] == A[i - 1] && (visited & (1 << (i - 1))) == 0) continue;
                  temp.add(A[i]);
                  dfs(A, temp, visited | (1 << i));
                  temp.remove(temp.size() - 1);
              }
          }
      }
      private boolean square(int sum){
          return (int)Math.sqrt(sum) * (int)Math.sqrt(sum) == sum;
      }
      }
    

Reference

1. 花花酱 LeetCode 996. Number of Squareful Arrays