Number of Squareful Arrays

Intuition

Solution

class Solution:
    def numSquarefulPerms(self, A: List[int]) -> int:
        counter = collections.Counter(A)
        self.squareful = set()
        for i in A:
            for j in A:
                s = i + j
                if int(s ** (1 / 2)) ** 2 == s:
                    self.squareful.add(s)

        return sum([self.dfs(i, counter) for i in set(A)])

    def dfs(self, vertex, counter):
        counter[vertex] -= 1
        ans = 0
        finished = True
        for child, count in counter.items():
            if count == 0:
                continue
            finished = False
            if child + vertex in self.squareful:
                ans += self.dfs(child, counter)
        counter[vertex] += 1
        if finished:
            return 1
        return ans