力扣—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

解题思路

本题就是需要找出一个能够连接图所有节点的哈密顿回路，对于一个列表，可以先使用Count类来存储各节点的数目，然后可以使用双重循环来遍历满足条件的顶点对，找到后两个顶点之间连接一条无向边，然后使用dfs来分别从每个节点来遍历，每次遍历的时候都要访问Count类中是否还有节点

import collections
class Solution:
    def __init__(self):
        self.sum = 0
    def numSquarefulPerms(self, A):
        n = len(A)
        count = collections.Counter(A)
        graph = {x: [] for x in A}
        for x in count:
            for y in count:
                if int((x + y) ** 0.5 + 0.5) ** 2 == x + y:
                    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
        for x in count:
            self.sum += dfs(x, len(A) - 1)
        return self.sum

S = Solution()
print(S.numSquarefulPerms([1, 17, 8]))