Number of Squareful Arrays LeetCode solution in Python

1
# Runtime: 56 ms (Top 51.37%) | Memory: 13.9 MB (Top 48.63%)
2
class Solution:
3
    def isSquare(self, num):
4
        return int(num**0.5) ** 2 == num
5

6
    def makePermutation(self, used, vis, prev, n):
7
        if used == n:
8
            # we reached the end
9
            self.ans += 1
10
            return
11
        tmp = {}
12
        for i in range(n):
13
            if vis[i] == False and self.nums[i] not in tmp:
14
                tmp[self.nums[i]] = True
15
                if self.nums[i] in self.d[prev]:
16
                    vis[i] = True
17
                    self.makePermutation(used + 1, vis, self.nums[i], n)
18
                    vis[i] = False
19

20
    def numSquarefulPerms(self, nums: List[int]) -> int:
21
        d = {x: {} for x in nums}
22
        n = len(nums)
23
        for i in range(n):
24
            for j in range(i + 1, n):
25
                if self.isSquare(nums[i] + nums[j]):
26
                    d[nums[i]][nums[j]] = True
27
                    d[nums[j]][nums[i]] = True
28
        self.nums = nums
29
        self.ans = 0
30
        self.d = d
31
        vis = [False] * n
32
        tmp = {}
33
        for i in range(n):
34
            if nums[i] in tmp:
35
                continue
36
            tmp[nums[i]] = True
37
            vis[i] = True
38
            self.makePermutation(1, vis, self.nums[i], n)
39
            vis[i] = False
40
        return self.ans