正方形数组的数目

标签: 位运算 数组 数学 动态规划 回溯 状态压缩

难度: Hard

给定一个非负整数数组 A,如果该数组每对相邻元素之和是一个完全平方数,则称这一数组为正方形数组。

返回 A 的正方形排列的数目。两个排列 A1A2 不同的充要条件是存在某个索引 i,使得 A1[i] != A2[i]。

示例 1:

输入:[1,17,8]
输出:2
解释:
[1,8,17] 和 [17,8,1] 都是有效的排列。

示例 2:

输入:[2,2,2]
输出:1

提示:

  1. 1 <= A.length <= 12
  2. 0 <= A[i] <= 1e9

Submission

运行时间: 22 ms

内存: 16.1 MB

class Solution:
    def numSquarefulPerms(self, nums: List[int]) -> int:
        n=len(nums)
        nums.sort()
        mask=0
        vis=[False]*n
        set1=set()
        def dfs(pre,mask):
            if mask==(1<<n)-1:
                return 1 
            ans=0
            for i in range(n):
                if (mask >>i) & 1:continue
                if i>0 and nums[i]==nums[i-1] and not vis[i-1]:continue
                num=nums[i]+pre
                if  mask==0 or(  isqrt(num)**2==num):
                    vis[i]=True
                    ans+=dfs(nums[i],mask |(1<<i))
                    vis[i]=False
            return ans
        ans=dfs(-1,mask)
        return ans
                 
            

Explain

这个题解采用了深度优先搜索(DFS) + 回溯 + 位掩码的方法。首先,对数组进行排序,以便后续处理重复元素。然后,使用一个位掩码来表示当前哪些元素已经被选中。对于每一个可能的元素,如果它和前一个元素相同,并且前一个元素没有被选中,则跳过,以避免重复排列。接着,检查当前选中的元素和前一个元素的和是否是完全平方数,如果是,则继续递归搜索下一个元素。当所有元素都被选中时,找到了一个有效的排列,返回1。最后,返回所有有效排列的总数。

时间复杂度: O(n!)

空间复杂度: O(n)

class Solution:
    def numSquarefulPerms(self, nums: List[int]) -> int:
        n = len(nums)
        nums.sort()
        vis = [False] * n
        def dfs(pre, mask):
            if mask == (1 << n) - 1:
                return 1
            ans = 0
            for i in range(n):
                if (mask >> i) & 1:
                    continue
                if i > 0 and nums[i] == nums[i - 1] and not vis[i - 1]:
                    continue
                num = nums[i] + pre
                if mask == 0 or (isqrt(num) ** 2 == num):
                    vis[i] = True
                    ans += dfs(nums[i], mask | (1 << i))
                    vis[i] = False
            return ans
        return dfs(-1, 0)

Explore

在题解中对数组进行排序的目的主要是为了方便处理重复元素,从而避免生成重复的排列。排序后的数组使得相同的元素聚集在一起,这样在递归搜索中可以通过简单的条件检查来跳过重复的元素,从而提高算法的效率。此外,排序也有助于在递归过程中更快地判断两个数的和是否为完全平方数,因为数组是有序的,可以更方便地选择和当前元素配对的候选数。

位掩码是一种使用二进制数字的每一位来表示一个独立变量的方法,通常用于跟踪多个布尔状态。在这个题解中,位掩码用来跟踪数组中每个元素是否已被选中。例如,如果有三个元素,位掩码'101'表示第一个和第三个元素被选中,而第二个元素未被选中。在递归函数中,通过位运算`mask >> i & 1`可以检查第i个元素是否已选中;通过`mask | (1 << i)`来更新状态,标记第i个元素为选中。这种方法能够有效且紧凑地跟踪每个元素的选中状态。

题解中的方法在大部分情况下是有效的,尤其是当数组经过排序之后。当数组中存在多个相同的元素时,排序保证了相同元素排列在一起。在遍历这些元素时,通过检查当前元素的前一个相同元素是否已被选中(即跳过条件`i > 0 and nums[i] == nums[i - 1] and not vis[i - 1]`),可以防止生成相同的排列。这种方法依赖于数组的排序状态以及递归过程中对元素选中状态的维护,是一种有效的去重策略。

`pre`变量在`dfs`递归函数中表示上一个被选择加入到当前排列中的元素的值。它用于与当前考虑加入排列的元素的值相加,以检查这两个数的和是否为完全平方数。如果是,递归继续,否则跳过当前元素。`pre`的使用是检查排列中相邻元素的和的关键,确保构建的排列满足题目条件,即排列中任意两个相邻数字的和都是完全平方数。函数流程和最终结果都依赖于`pre`的正确更新和使用,因此影响着搜索过程中的决策和排列的有效性。