# 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].
#
#
#
# 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
import collections
import math


class Solution(object):
    def numSquarefulPerms(self, A):
        """
        :type A: List[int]
        :rtype: int
        """

        def isSquare(x):
            return math.sqrt(x) == int(math.sqrt(x))

        graph = collections.defaultdict(set)

        for i in xrange(len(A)):
            for j in xrange(i):
                if isSquare(A[i] + A[j]):
                    graph[i].add(j)
                    graph[j].add(i)

        self.res = 0

        for i in xrange(len(A)):
            if i > 0 and A[i] == A[i - 1]:
                continue
            self.dfs(A, graph, i, 0, set())

        return self.res

    def dfs(self, A, graph, i, level, visited):
        visited.add(i)
        if level == len(A) - 1:
            self.res += 1
            return
        for k in graph[i]:
            if k not in visited:
                self.dfs(A, graph, k, level + 1, visited)
        return

    def numSquarefulPerms2(self, A):
        """
        :type A: List[int]
        :rtype: int
        """

        def isSquare(x):
            return math.sqrt(x) == int(math.sqrt(x))

        graph = collections.defaultdict(set)

        for i in xrange(len(A)):
            for j in xrange(i):
                if isSquare(A[i] + A[j]):
                    graph[i].add(j)
                    graph[j].add(i)

        self.res, visited = 0, []

        for i in xrange(len(A)):
            if i > 0 and A[i] == A[i - 1]:
                continue

            # visited.append(i)
            self.dfs2(A, graph, i, 0, [])
            # visited.pop()

        return self.res

    def dfs2(self, A, graph, i, level, visited):
        visited.append(i)
        if level == len(A) - 1:
            self.res += 1
            print visited
            return
        for k in graph[i]:
            if k not in visited:
                # visited.append(k)
                self.dfs2(A, graph, k, level + 1, visited)
                # visited.pop()
        visited.pop()
        return

    def numSquarefulPerms3(self, A):
        """
        :type A: List[int]
        :rtype: int
        """

        def isSquare(x):
            return math.sqrt(x) == int(math.sqrt(x))

        def dfs(i, level):
            count[i] -= 1
            if not level:
                # bugfixed we can not return here, since this is a backtrack
                # if you directly return, then you are losing the opportunity to do `count[i]+=1`
                # return 1
                ans = 1
            else:
                ans = 0
                for j in graph[i]:
                    if count[j]:
                        ans += dfs(j, level - 1)
            count[i] += 1
            return ans

        graph = collections.defaultdict(set)

        # notice for test case[2,2,2], you woudl get [2:3]
        count = collections.Counter(A)
        # we are looping counter's key, we use i, j to loop all keys
        # but if we use the case twice , what would happend, for ex [2, 14]
        # we got [2:1, 14:1], then we would have graph[2] = {2, 14}, isn't this wrong?
        # but in dfs function, we have check, `if count[j]`, assume we used 2, then count[2] = 0
        # we will not be able to use `2` again. so no need to worry
        for i in count:
            for j in count:
                # add this check, the time is from 20 -> 28 ms, from 100% -> 32%, as above comments said
                # since in dfs function, we have teh check, so we don't need to worry.
                # if i == j and count[i] == 1:
                #     continue
                if isSquare(i + j):
                    # for test case [2,2,2], you got graph[2] = {2}
                    graph[i].add(j)

        return sum(dfs(i, len(A) - 1) for i in count)



if __name__ == '__main__':
    print Solution().numSquarefulPerms2([0, 1, 0, 1])
    print Solution().numSquarefulPerms2([0, 0, 1, 1])
    print Solution().numSquarefulPerms([0, 1, 0, 1])
    print Solution().numSquarefulPerms([0, 1, 0, 1, 0, 1])
    print Solution().numSquarefulPerms([0, 0, 0, 1, 1, 1])
    print Solution().numSquarefulPerms3([2, 2, 2])
    print Solution().numSquarefulPerms([1, 17, 8])
