996. Number of Squareful Arrays

Problem

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

Problem Summary

Given an array A of non-negative integers, if the sum of every pair of adjacent elements in the array is a perfect square, then the array is called a squareful array.

Return the number of squareful permutations of A. Two permutations A1 and A2 are different if and only if there exists some index i such that A1[i] != A2[i].

Solution Ideas

This problem is an enhanced version of Problem 47. Problem 47 asks for all unique permutations of an array. This problem asks for all unique permutations where the sum of every two adjacent numbers is a perfect square.
The idea is exactly the same as Problem 47, except that an additional check is added to determine whether the sum of two adjacent numbers is a perfect square. Note that during the DFS process, pruning is needed; otherwise the time complexity will be very high and cause a timeout.

Code


package leetcode

import (
	"math"
	"sort"
)

func numSquarefulPerms(A []int) int {
	if len(A) == 0 {
		return 0
	}
	used, p, res := make([]bool, len(A)), []int{}, [][]int{}
	sort.Ints(A) // This is the key logic for deduplication
	generatePermutation996(A, 0, p, &res, &used)
	return len(res)
}

func generatePermutation996(nums []int, index int, p []int, res *[][]int, used *[]bool) {
	if index == len(nums) {
		checkSquareful := true
		for i := 0; i < len(p)-1; i++ {
			if !checkSquare(p[i] + p[i+1]) {
				checkSquareful = false
				break
			}
		}
		if checkSquareful {
			temp := make([]int, len(p))
			copy(temp, p)
			*res = append(*res, temp)
		}
		return
	}
	for i := 0; i < len(nums); i++ {
		if !(*used)[i] {
			if i > 0 && nums[i] == nums[i-1] && !(*used)[i-1] { // This is the key logic for deduplication
				continue
			}
			if len(p) > 0 && !checkSquare(nums[i]+p[len(p)-1]) { // Key pruning condition
				continue
			}
			(*used)[i] = true
			p = append(p, nums[i])
			generatePermutation996(nums, index+1, p, res, used)
			p = p[:len(p)-1]
			(*used)[i] = false
		}
	}
	return
}

func checkSquare(num int) bool {
	tmp := math.Sqrt(float64(num))
	if int(tmp)*int(tmp) == num {
		return true
	}
	return false
}