LeetCode：996. Number of Squareful Arrays - Python

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

问题分析：

题号996在互联网公司是多么吉利的一个数字哈。说说题目，一开始就想到了回溯法，递归求解，后来看了官方的解答，也是回溯法和动规法，而且是时间复杂度都很高，基本接近穷举法，这就让人感觉这个题目很无聊了。现在参考官方解答总结一下使用回溯法的解决方案：

（1）可以先统计一下每个数字的出现的次数，并且统计一下，每个数字与其他能组成平方数的数字放到一块，这样做可以节省一定的计算量。

（2）使用回溯法，依次选出一个数字进行深度优先 递归判断。

Python3实现：

# 官方解答
import collections


class Solution():
    def numSquarefulPerms(self, A):
        count = collections.Counter(A)

        graph = {x: [] for x in count}
        for x in count:  # 求出那些数字是和x 可以组成平方数的
            for y in count:
                if int((x+y)**.5 + 0.5) ** 2 == x+y:  # 测试是，0.5 加与不加 不影响结果
                    graph[x].append(y)

        def dfs(x, todo):
            count[x] -= 1
            if todo == 0:
                ans = 1
            else:
                ans = 0
                for y in graph[x]:  # 依次选出一个数字进行深度优先 递归判断
                    if count[y]:
                        ans += dfs(y, todo - 1)
            count[x] += 1
            return ans

        return sum(dfs(x, len(A) - 1) for x in count)


if __name__ == '__main__':
    solu = Solution()
    print(solu.numSquarefulPerms([1, 17, 8]))

声明： 总结学习，有问题或不当之处，可以批评指正哦，谢谢。

题目链接、参考链接