2597.美丽子集的数目

‘美丽子集的数目’题解借助 `Counter` 统计元素次数，筛选出公差为 `k` 的等差数列首项，用位运算和动态规划计算各数列满足元素差值条件的子集方案数，依排列组合原理相乘并排除空集得美丽子集数量。
  • 本文作者
  • 文章发布日期 2025-03-07 19:00
  • 热度 38
  • 作者心情 非常愉快
  • 本文共计
  • 预计阅读

题目要求

给你一个由正整数组成的数组 nums 和一个 正 整数 k 。

如果 nums 的子集中，任意两个整数的绝对差均不等于 k ，则认为该子数组是一个 美丽 子集。

返回数组 nums 中 非空 且 美丽 的子集数目。

nums 的子集定义为：可以经由 nums 删除某些元素（也可能不删除）得到的一个数组。只有在删除元素时选择的索引不同的情况下，两个子集才会被视作是不同的子集。

示例 1：

输入：nums = [2,4,6], k = 2
输出：4
解释：数组 nums 中的美丽子集有：[2], [4], [6], [2, 6] 。
可以证明数组 [2,4,6] 中只存在 4 个美丽子集。

示例 2：

输入：nums = [1], k = 1
输出：1
解释：数组 nums 中的美丽数组有：[1] 。
可以证明数组 [1] 中只存在 1 个美丽子集。

提示：

  • 1 <= nums.length <= 18

  • 1 <= nums[i], k <= 1000

题目解析参考

参考答案

class Solution:
    def beautifulSubsets(self, nums: List[int], k: int) -> int:
        cnt = Counter(nums)
        ans = 1
        for x, c in cnt.items():
            if x - k in cnt:  # x 不是等差数列的首项
                continue
            # 计算这一组的方案数
            f0, f1 = 1, 1 << c
            x += k
            while x in cnt:
                f0, f1 = f1, f1 + f0 * ((1 << cnt[x]) - 1)
                x += k
            ans *= f1  # 每组方案数相乘
        return ans - 1  # 去掉空集

代码整体思路

  1. 统计数组中每个元素的出现次数。

  2. 找出所有以某个元素为首项、公差为 k 的等差数列。

  3. 对每个等差数列，使用动态规划的方法计算满足条件的子集方案数。

  4. 将所有等差数列的方案数相乘得到最终结果。

  5. 减去空集的情况，得到最终的美丽子集数量。

代码详细解读

统计元素出现次数

cnt = Counter(nums)

  • Counter 是 Python 的 collections 模块中的一个类，它可以方便地统计可迭代对象中每个元素的出现次数。在这里，我们将数组 nums 作为参数传入 Counter 类，创建了一个名为 cnt 的字典。字典的键是数组 nums 中的元素，对应的值是该元素在数组中出现的次数。

  • 例如，如果 nums = [1, 2, 2, 3]，那么 cnt 将会是 {1: 1, 2: 2, 3: 1}。

初始化最终结果

ans = 1

  • 我们将最终结果 ans 初始化为 1。这是因为后续我们要把每个等差数列组的方案数相乘，初始值为 1 不会影响最终的乘积结果。乘法运算中，任何数乘以 1 都等于其本身。

遍历元素并筛选等差数列首项

for x, c in cnt.items():
    if x - k in cnt:  # x 不是等差数列的首项
        continue

  • cnt.items() 会遍历 cnt 字典中的每一个键值对。其中，x 代表当前元素的值，c 代表该元素出现的次数。

  • if x - k in cnt: 这一条件判断用于确定当前元素 x 是否是一个以公差为 k 的等差数列的首项。如果 x - k 也在 cnt 字典中，说明存在比 x 更小的元素可以作为这个等差数列的首项，所以 x 不是首项，我们使用 continue 语句跳过当前循环，直接处理下一个元素。

  • 例如，当 k = 1，cnt = {1: 1, 2: 1, 3: 1} 时，如果当前 x = 2，由于 2 - 1 = 1 也在 cnt 中，所以 2 不是首项，会跳过对 2 的处理。

计算以当前元素为首项的等差数列组的方案数

f0, f1 = 1, 1 << c
x += k
while x in cnt:
    f0, f1 = f1, f1 + f0 * ((1 << cnt[x]) - 1)
    x += k

  • 初始化方案数：

    • f0 = 1：f0 表示不选择当前这组等差数列元素的方案数，也就是空集的情况。因为空集是唯一的一种不包含任何元素的集合，所以方案数为 1。

    • f1 = 1 << c：这里涉及到位运算中的左移操作。左移运算符 << 可以快速计算 2 的幂次方，1 << c 等价于 2^c。c 是当前首项 x 的出现次数，首项 x 可以出现 0 次、1 次、……、c 次，总共 2^c 种不同的选择情况，所以 f1 表示只考虑首项 x 时的方案数。

    • 例如，如果首项 x = 1 出现了 2 次（即 c = 2），那么 1 << 2 等于 4，意味着首项 1 有 4 种选择情况：出现 0 次、出现 1 次、出现 2 次。

  • 动态规划更新方案数：

    • x += k：将 x 更新为当前等差数列的下一个元素。因为我们要处理以 x 为首项、公差为 k 的等差数列，所以每次将 x 加上 k 就可以得到下一个元素。

    • while x in cnt:：这个 while 循环会不断检查更新后的 x 是否在 cnt 字典中。如果在，说明这个等差数列还在继续，我们需要继续计算方案数；如果不在，说明等差数列结束，退出循环。

    • f0, f1 = f1, f1 + f0 * ((1 << cnt[x]) - 1)：这是动态规划的核心状态转移方程，用于更新方案数。

      • f0 更新为 f1：在下一步计算中，f0 要代表上一轮的 f1，即不选择当前元素时的方案数。

      • (1 << cnt[x]) - 1：1 << cnt[x] 表示当前元素 x 所有可能的出现情况（包括出现 0 次），减去 1 是排除当前元素 x 出现 0 次的情况，只考虑它至少出现 1 次的情况。

      • f0 * ((1 << cnt[x]) - 1)：表示选择当前元素 x 的方案数。因为 f0 是不选择当前元素时的方案数，当我们选择当前元素 x 时，它至少出现 1 次有 (1 << cnt[x]) - 1 种情况，两者相乘就是选择当前元素 x 的方案数。

      • f1 更新为 f1 + f0 * ((1 << cnt[x]) - 1)：f1 原本表示不选当前元素 x 的方案数，加上选择当前元素 x 的方案数，就得到了考虑当前元素 x 时的总方案数。

    • x += k：继续将 x 更新为等差数列的下一个元素，以便继续循环处理。

汇总每组方案数

ans *= f1  # 每组方案数相乘

  • 每完成一组以某个元素为首项的等差数列的方案数计算（得到 f1）后，我们将其乘到最终结果 ans 中。这是因为不同组的等差数列之间不会相互影响，根据排列组合的乘法原理，将它们各自的方案数相乘就可以得到所有可能的组合方案数。

排除空集情况

return ans - 1  # 去掉空集

  • 题目要求的是除了空集之外的美丽子集数量。因为我们在计算过程中把空集的情况也包含进去了（例如 f0 = 1 就代表了空集的方案数），所以最后要从总方案数 ans 中减去空集这一种情况，得到最终的美丽子集数量。

复杂度分析

  • 时间复杂度：代码的主要时间开销在于遍历 cnt 字典以及处理每个等差数列。假设数组 nums 的长度为 n，不同元素的数量为 m，那么时间复杂度大致为 \(O(m)\)，因为每个元素最多被处理一次。

  • 空间复杂度：主要的空间开销是 cnt 字典，它存储了数组中每个元素的出现次数，所以空间复杂度为 \(O(m)\)，其中 m 是不同元素的数量。