等差数列划分 II - 子序列

标签: 数组 动态规划

难度: Hard

给你一个整数数组 nums ，返回 nums 中所有 等差子序列 的数目。

如果一个序列中 至少有三个元素 ，并且任意两个相邻元素之差相同，则称该序列为等差序列。

  • 例如，[1, 3, 5, 7, 9]、[7, 7, 7, 7] 和 [3, -1, -5, -9] 都是等差序列。
  • 再例如，[1, 1, 2, 5, 7] 不是等差序列。

数组中的子序列是从数组中删除一些元素（也可能不删除）得到的一个序列。

  • 例如，[2,5,10] 是 [1,2,1,2,4,1,5,10] 的一个子序列。

题目数据保证答案是一个 32-bit 整数。

示例 1：

输入：nums = [2,4,6,8,10]
输出：7
解释：所有的等差子序列为：
[2,4,6]
[4,6,8]
[6,8,10]
[2,4,6,8]
[4,6,8,10]
[2,4,6,8,10]
[2,6,10]

示例 2：

输入：nums = [7,7,7,7,7]
输出：16
解释：数组中的任意子序列都是等差子序列。

提示：

  • 1  <= nums.length <= 1000
  • -231 <= nums[i] <= 231 - 1

Submission

运行时间: 259 ms

内存: 0.0 MB

class Solution:
    def numberOfArithmeticSlices(self, nums: List[int]) -> int:
        N, ans = len(nums), 0
        dp, left, right = [defaultdict(int) for _ in range(N)], defaultdict(int), defaultdict(list)
        for i in range(N - 1, -1, -1):
            right[nums[i]].append(i)
        for i, m in enumerate(nums):
            right[m].pop()
            if not right[m]: del right[m]
            lr = ((l, m*2 - l) for l in left if m*2 - l in right) if len(left) < len(right) else ((m*2 - r, r) for r in right if m*2 - r in left)
            for l, r in lr:
                cnt = left[l] + dp[i][l]
                for j in right[r]:
                    dp[j][m] += cnt
            ans += sum(dp[i].values())
            left[m] += 1
        return ans

Explain

这个题解使用动态规划的方法来解决问题。主要思路如下： 1. 使用两个字典 left 和 right 分别记录每个数字最左边和最右边的出现位置。 2. 从右往左遍历数组，对于每个位置 i 上的数字 m： - 更新 right[m]，删除当前位置 - 枚举所有可能的等差数列的另外两个数 l 和 r，其中 m = (l + r) / 2 - 对于每个符合条件的 l 和 r，累加 dp[i][l] 的值到 dp[j][m] 上，其中 j 为 right[r] 中的每个位置 - 将 dp[i] 的所有值累加到答案 ans 上 - 将当前位置 i 加入到 left[m] 中 3. 最终返回 ans 即可

时间复杂度: O(n^2)

空间复杂度: O(n)

class Solution:
    def numberOfArithmeticSlices(self, nums: List[int]) -> int:
        N, ans = len(nums), 0
        # dp[i][j] 表示以 nums[i] 结尾、公差为 j 的等差数列的个数
        dp, left, right = [defaultdict(int) for _ in range(N)], defaultdict(int), defaultdict(list) 
        for i in range(N - 1, -1, -1):
            right[nums[i]].append(i)
        for i, m in enumerate(nums):
            right[m].pop()
            if not right[m]: del right[m]
            # 枚举所有可能的等差数列的另外两个数 l 和 r
            lr = ((l, m*2 - l) for l in left if m*2 - l in right) if len(left) < len(right) else ((m*2 - r, r) for r in right if m*2 - r in left)
            for l, r in lr:
                cnt = left[l] + dp[i][l]
                for j in right[r]:
                    dp[j][m] += cnt
            ans += sum(dp[i].values())
            left[m] += 1
        return ans

Explore

从右向左遍历数组允许我们在处理每个元素时，已经有了该元素右侧所有元素的信息，这对于动态规划中状态的更新非常关键。通过这种方式，我们可以确保当我们计算以 nums[i] 结尾的等差数列个数时，其右侧的数列个数已经被计算过，从而可以直接使用这些信息来更新当前位置的状态。这样的遍历顺序帮助我们避免了重复计算并保证了每个状态是基于最终的、更新过的值。

在这个动态规划解法中，dp[i][j] 中的 j 实际上代表的是等差数列的第一个数，而不是常见的公差。dp[i][j] 表示以 nums[i] 结尾，且以 j 作为等差数列第一个数时的等差数列的个数。更新 dp 数组时，我们需要检查所有可能的数对 (l, r) 使得 nums[i] = (l + r) / 2。对于每个符合条件的 l 和 r，我们更新 dp[j][m] 的值，其中 j 是 r 在数组中的位置，m 是当前的 nums[i]。这样，dp[j][m] 的值就是在位置 j 以 m 作为等差数列倒数第二个数时的数列个数。

在生成数对 (l, r) 时同时考虑 left 和 right 字典可以优化算法的性能。这种方法通过考虑字典中元素较少的一方来减少需要检查的数对的数量。如果 left 中的元素比 right 少，我们从 left 生成数对，反之则从 right。这样可以减少循环的次数，因为我们只遍历较小的字典并检查另一个字典是否含有符合条件的元素，从而提高算法的效率。

在这个算法中，left 和 right 字典分别用来记录数组中每个数字出现的位置。right 字典在遍历数组时首先被填充，记录每个数字出现的所有位置。随着从右向左的遍历，每到一个新位置，该位置的索引从 right 字典对应的列表中删除。如果某个数字的索引列表变空，则从字典中完全移除该数字的条目。left 字典在遍历过程中逐步构建，记录每个数字最新的出现位置。这样，两个字典共同支持了算法中对于等差数列两端点的快速查找和状态更新。