996. Number of Squareful Arrays

题目：

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

996. 正方形数组的数目

给定一个非负整数数组 A，如果该数组每对相邻元素之和是一个完全平方数，则称这一数组为正方形数组。

返回 A 的正方形排列的数目。两个排列 A1 和 A2 不同的充要条件是存在某个索引 i，使得 A1[i] != A2[i]。

示例 1：

输入：[1,17,8]
输出：2
解释：
[1,8,17] 和 [17,8,1] 都是有效的排列。

示例 2：

输入：[2,2,2]
输出：1

提示：

1. 1 <= A.length <= 12
2. 0 <= A[i] <= 1e9

思路：

背包问题，深度优先搜索

代码：

class Solution:
    def numSquarefulPerms(self, A: List[int]) -> int:
        from collections import Counter
        self.c = Counter(A)
        self.cond = {i: {j for j in self.c if int((i+j)**0.5)**2 == i+j} for i in self.c}
        self.cnt = 0

        def dfs(x, left=len(A)-1):
            self.c[x] -= 1
            if left == 0:
                self.cnt += 1
            for y in self.cond[x]:
                if self.c[y]:
                    dfs(y, left-1)
            self.c[x] += 1

        for x in self.c:
            dfs(x)
        return self.cnt

import "math"

func isSquare(x int) bool {
    r := int(math.Sqrt(float64(x)))
    return r*r == x
}

func numSqurefulPerms(A []int) int {
    size := len(A)
    count := make(map[int]int, size)
    for _, a := range A {
        count[a]++
    }
    conds := make(map[int][]int, size)
    for x := range count {
        for y := range count {
            if isSquare(x + y) {
                conds[x] = append(conds[x], y)
            }
        }
    }

    cnt := 0
    var dfs func(int, int)

    dfs = func(x, left int) {
        if left == 0 {
            cnt++
            return
        }
        count[x]--
        for _, y := range conds[x] {
            if count[y] > 0 {
                dfs(y, left-1)
            }
        }
        count[x]++
    }

    for x := range count {
        dfs(x, size-1)
    }
    return cnt

}