Skip to content

996. Number of Squareful Arrays

Difficulty Topics

Description

An array is squareful if the sum of every pair of adjacent elements is a perfect square.

Given an integer array nums, return the number of permutations of nums that are squareful.

Two permutations perm1 and perm2 are different if there is some index i such that perm1[i] != perm2[i].

 

Example 1:

Input: nums = [1,17,8]
Output: 2
Explanation: [1,8,17] and [17,8,1] are the valid permutations.

Example 2:

Input: nums = [2,2,2]
Output: 1

 

Constraints:

  • 1 <= nums.length <= 12
  • 0 <= nums[i] <= 109

Solution

number-of-squareful-arrays.py
class Solution:
    def numSquarefulPerms(self, nums: List[int]) -> int:
        n = len(nums)
        nums.sort()
        res = 0
        used = [False] * n

        @cache
        def good(m):
            sq = int(math.sqrt(m))

            return sq * sq == m

        def perm(path):
            nonlocal res

            if len(path) == n:
                res += 1
                return

            for i in range(n):
                if used[i]: continue

                if i > 0 and nums[i] == nums[i - 1] and not used[i - 1]: continue

                if len(path) == 0 or (len(path) > 0 and good(path[-1] + nums[i])):
                    used[i] = True
                    perm(path + [nums[i]])
                    used[i] = False

        perm([])        

        return res