996. Number of Squareful Arrays (Hard) - TengnanYao/daily_leetcode GitHub Wiki

class Solution(object):
    def numSquarefulPerms(self, nums):
        """
        :type nums: List[int]
        :rtype: int
        """
        def check(num):
            l, r = 0, num
            while l <= r:
                m = (l + r) // 2
                if m * m == num:
                    return True
                if m * m > num:
                    r = m - 1
                else:
                    l = m + 1
            return False
        def dfs(nums, pre):
            if not nums:
                self.result += 1
            else:
                seen = set()
                for i, num in enumerate(nums):
                    if check(num + pre) and num not in seen:
                        seen.add(num)
                        dfs(nums[ : i] + nums[i + 1 : ], num)
        self.result = 0
        s = set(nums)
        for num in s:
            temp = nums[ : ]
            temp.remove(num)
            dfs(temp, num)
        return self.result