【LeetCode 446】等差数列划分 II - 子序列

• 链接：https://leetcode.cn/problems/arithmetic-slices-ii-subsequence/description/

难点

#

我们首先考虑至少有两个元素的等差子序列，下文将其称作弱等差子序列。

由于尾项和公差可以确定一个等差数列，因此我们定义状态 dp[i][d] 表示尾项为 nums[i]，公差为 d 的弱等差子序列的个数。

我们用一个二重循环去遍历 nums 中的所有元素对 (nums[i], nums[j])，其中 j<i。将 nums[i] 和 nums[j] 分别当作等差数列的尾项和倒数第二项，则该等差数列的公差 d=nums[i]-nums[j]。由于公差相同，我们可以将 nums[i] 加到以 nums[j] 为尾项，公差为 d 的弱等差子序列的末尾，这对应着状态转移 dp[i][d]+=dp[j][d]。同时，(nums[i], nums[j]) 这一对元素也可以当作一个弱等差子序列，故有状态转移dp[i][d]+=dp[j][d]+1

由于题目要统计的等差子序列至少有三个元素，我们回顾上述二重循环，其中「将 nums[i] 加到以 nums[j] 为尾项，公差为 d 的弱等差子序列的末尾」这一操作，实际上就构成了一个至少有三个元素的等差子序列，因此我们将循环中的 dp[j][d] 累加，即为答案。

代码

#

class Solution {
    public int numberOfArithmeticSlices(int[] nums) {
        int n = nums.length;
        Map<Long, Integer>[] dp = new HashMap[n];
        for (int i = 0; i < n; i++) {
            dp[i] = new HashMap<>();
        }
        int ans = 0;
        for(int i = 1; i < n; i++) {
            for(int j = 0; j < i; j++) {
                long d = 1L * nums[i] - nums[j];
                int temp = dp[j].getOrDefault(d, 0);
                ans += temp;
                dp[i].put(d, temp + 1 + dp[i].getOrDefault(d, 0));
            }
        }
        return ans;
    }
}