996. Number of Squareful Arrays #
题目 #
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 <= 120 <= A[i] <= 1e9
题目大意 #
给定一个非负整数数组 A，如果该数组每对相邻元素之和是一个完全平方数，则称这一数组为正方形数组。
返回 A 的正方形排列的数目。两个排列 A1 和 A2 不同的充要条件是存在某个索引 i，使得 A1[i] != A2[i]。
解题思路 #
- 这一题是第 47 题的加强版。第 47 题要求求出一个数组的所有不重复的排列。这一题要求求出一个数组的所有不重复，且相邻两个数字之和都为完全平方数的排列。
- 思路和第 47 题完全一致，只不过增加判断相邻两个数字之和为完全平方数的判断，注意在 DFS 的过程中，需要剪枝，否则时间复杂度很高，会超时。
代码 #
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) // 这里是去重的关键逻辑
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] { // 这里是去重的关键逻辑
continue
}
if len(p) > 0 && !checkSquare(nums[i]+p[len(p)-1]) { // 关键的剪枝条件
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
}