题目描述

给你一个整数数组 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 整数。

样例

输入：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]

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

限制

• 1 <= nums.length <= 1000
• -2^31 <= nums[i] <= 2^31 - 1

算法

(动态规划) $O(n^2)$

1. 设 $f(i, d)$ 表示以 $i$ 为结尾的差值为 $d$ 的子序列的个数（包括长度为 2 的）。
2. 初始时 $f(i, d)$ 都是 $0$。转移时，枚举 $0 \le j < i$，令 $d = nums[i] - nums[j]$，则 $f(i, d) += f(j, d) + 1$，这个的意思是以 $j$ 结尾的子序列都可以接到 $nums[i]$ 上，nums[i], nums[j] 也算 1 个。
3. 最终答案为 $\sum_i \sum_d f(i, d)$，然后减去 $n(n-1)/2$。

时间复杂度

• 一共有 $O(n^2)$ 个状态，每次转移复杂度为 $O(1)$，故总时间复杂度为 $O(n^2)$。

空间复杂度

• 需要 $O(n^2)$ 个哈希表来存储状态。故空间复杂度为 $O(n^2)$。

C++ 代码

#define LL long long

class Solution {
public:
    int numberOfArithmeticSlices(vector<int>& nums) {
        const int n = nums.size();
        vector<unordered_map<LL, int>> f(n);

        int ans = 0;
        for (int i = 1; i < n; i++)
            for (int j = 0; j < i; j++) {
                LL d = (LL)(nums[i]) - nums[j];
                int sum = 0;
                if (f[j].find(d) != f[j].end())
                    sum = f[j][d];

                f[i][d] += sum + 1;
                ans += sum + 1;
            }

        return ans - n * (n - 1) / 2;
    }
};