LeetCode996.Number of Sequareful Arrays(正方形数组的数目)

Description

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].

---

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

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

题目链接：https://leetcode.com/problems/number-of-squareful-arrays/
个人网站：http://redtongue.cn or http://redtongue.github.io/

Difficulty: hard

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

分析

• Solution

参考代码

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