Leetcode 996 Number of Squareful Arrays

发表于 |

给定一个非负整数数组 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

思路分析：
我们就使用最笨的办法进行分析，先随便放一个到数组中的第一个位置，那么第二个位置就只能放几个值了，因为要满足相加为平方数这个条件。同样，第三个位置也只能放某几个数了。

那我们就将所有的数字看做节点，能相邻放的数字（也就是相加为平方数）用线相连，那么这样就构造出了一个图。

在这个图中，我们使用dfs方法进行遍历，由于A的长度不超过12，所以直接遍历不会超时。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
class Solution:
    def numSquarefulPerms(self, A: 'List[int]') -> 'int':
        d = collections.defaultdict(int)
        n = len(A)
        for i in range(n):
            for j in range(i+1, n):
                if int((A[i]+A[j])**0.5)**2 == A[i] + A[j]:
                    d[A[i]] |= 1 << j
                    d[A[j]] |= 1 << i
    
        self.res = 0
        vis = set()
        def dfs(num, c):
            # once all element used, means we find a permutations
            if all(c[key]==0 for key in c):
                self.res += 1
            n_vis = set() # avoid repeat element, such as [2,2,2]
            for j in range(32):
                mask = 1 << j
                # if A[j] can be next element and A[j] is not all used already
                if d[num] & mask and c[A[j]] > 0 and A[j] not in n_vis:
                    n_vis.add(A[j])
                    c[A[j]] -= 1 # use one, so reduce one
                    dfs(A[j], c)
                    c[A[j]] += 1 # after dfs, plus one
        
        for i in range(n):
            if A[i] in vis: continue
            vis.add(A[i])
            c = collections.Counter(A) # record the number of elements we can still use 
            c[A[i]] -= 1
            dfs(A[i], c)
        return self.res