动态规划的应用（五）：LeetCode 413, 446. 等差数列划分

LeetCode 413. Arithmetic Slices（等差数列划分）
LeetCode 446. Arithmetic Slices II - Subsequence（等差数列划分 II - 子序列）

其他动态规划的应用实例：
动态规划的应用（一）：最短路问题
动态规划的应用（二）：cutting stock 问题
动态规划的应用（三）：字符串相关问题
动态规划的应用（四）：LeetCode 1900. 最佳运动员的比拼回合
动态规划的应用（六）：矩阵相关问题

413. Arithmetic Slices（等差数列划分）

问题描述

思路与代码

动态规划：存储中间结果

看到题自然地想到了动态规划。
由于在迭代过程中，每增加一个元素，如果该元素与前两项组成的子数列为等差数列，那么之前找到的所有含有前两项的等差子数列均会产生一个新等差子数列。因此笔者定义了一个列表，以存储所有找到的等差子数列以便于操作：

class Solution:
    def numberOfArithmeticSlices(self, nums: List[int]) -> int:
        def dp(array: List[int]) -> List[List[int]]:
            """
            动态规划的递归函数
            :param array:  数列
            :return:  等差子数列的列表
            """

            # special case 1: 长度小于 3 时，返回空
            if len(array) < 3:
                return []

            # special case 2: 长度为 3 时，判断是否为等差数列
            elif len(array) == 3:
                if array[2] - array[1] == array[1] - array[0]:
                    return [[0, 1, 2]]
                else:
                    return []

            else:
                list_sub = dp(array=array[: -1])
                list_sub_add = []
                if array[-1] - array[-2] == array[-2] - array[-3]:  # 前一数列的最末两项与新元素形成等差数列
                    list_sub_add.append([len(array) - 3, len(array) - 2, len(array) - 1])

                    # 所有含有前一数列最末两项的等差子数列均可延长一项
                    for sub in list_sub:
                        if len(array) - 3 in sub and len(array) - 2 in sub:
                            list_sub_add.append(sub + [len(array) - 1])
                list_sub += list_sub_add

                return list_sub

        list_sub_ = dp(array=nums)

        return len(list_sub_)

代码运行通过，但是运行效果并不理想：

动态规划：不存储中间结果

参考了官方题解，发现了问题所在，即在迭代过程中，虽然新元素产生的新等差子数列与前两项元素有关，但之前找到的含前两项的等差子数列其实是在上一次迭代中找到的。因此仅从等差子数列的数量的角度上讲，每次迭代的新增数量之间存在一个迭代关系。

LeetCode 413 官方题解

代码如下：

class Solution:
    def numberOfArithmeticSlices(self, nums: List[int]) -> int:
        n = len(nums)
        if n == 1:
            return 0
        
        d, t = nums[0] - nums[1], 0
        ans = 0
        
        # 因为等差数列的长度至少为 3，所以可以从 i=2 开始枚举
        for i in range(2, n):
            if nums[i - 1] - nums[i] == d:
                t += 1
            else:
                d = nums[i - 1] - nums[i]
                t = 0
            ans += t
        
        return ans

运行效果：

非动态规划解法

此外，还有一种非官方解法，并没有使用动态规划，时间复杂度为 $O(n)$：

LeetCode 413 题解 非动态规划

代码如下：

class Solution:
    def numberOfArithmeticSlices(self, nums: List[int]) -> int:
        
        # 第一次遍历
        diffs = []
        for i in range(len(nums) - 1):
            diffs.append(nums[i + 1] - nums[i])
            
        # 第二次遍历
        cons = []
        a = 1
        for i in range(1, len(diffs)):
            if diffs[i] == diffs[i - 1]:
                a += 1
            else:
                cons.append(a)
                a = 1
        cons.append(a)
        
        # 第三次遍历
        res = 0
        for num in cons:
            res += int(self.calc(num))
        return res
        
    # 用于计算cons内每个数代表的等差数列个数
    def calc(self, n):
        if n == 1:
            return 0
        n += 1
        return (n - 2) * (n - 1) / 2

运行效果：

446. Arithmetic Slices II - Subsequence（等差数列划分 II - 子序列）

问题描述

思路与代码

本题与前一题（LeetCode 413）的区别在于，本题的等差子序列允许元素不连续。

受前一题的思路启发，针对等差子序列的元素可以不连续的情况，笔者定义了一个二维数组（矩阵），用来存储以第 $i$ 个元素和第 $j$ 个元素收尾的等差子序列的数量。在迭代过程中，如果发现新元素与前面某两个元素构成等差数列，则增加的等差子序列数量为，前面找到的以这两个元素收尾的等差子序列数 $+1$。

因此，基于动态规划的思路，可写出代码如下：

class Solution:
    def numberOfArithmeticSlices(self, nums: List[int]) -> int:
        n = len(nums)

        if n < 3:
            return 0

        ans = 0
        mat_tail = [[0 for _ in range(n)] for _ in range(n)]  # 以某两个元素收尾的等差子序列的个数
        for k in range(2, n):
            for i in range(k - 1):
                for j in range(i + 1, k):
                    if nums[k] - nums[j] == nums[j] - nums[i]:
                        num_add = mat_tail[i][j] + 1
                        ans += num_add
                        mat_tail[j][k] += num_add

        return ans

显然，该算法的时间复杂度为 $O(n^3)$。

然而，提交运行时发现，101 个算例中，第 99 个运行超时：

该算例的长度为 1000，看起来前面的代码的复杂度还是有点高……

参考官方题解，发现了问题所在，其实这里的动态规划是不需要 3 层循环的，只需要两层循环即可，具体的做法是把前述二维矩阵的第二维改为字典（哈希表），并将第二维从元素序号修改为第一维元素的前向差值。

LeetCode 446 官方题解

代码如下：

class Solution:
    def numberOfArithmeticSlices(self, nums: List[int]) -> int:
        ans = 0
        f = [defaultdict(int) for _ in nums]
        for i, x in enumerate(nums):
            for j in range(i):
                d = x - nums[j]
                cnt = f[j][d]
                ans += cnt
                f[i][d] += cnt + 1
        return ans

其中，调用 defaultdict 方法在本地运行时，需要：

from collections import defaultdict

优化后的运行效果：