Description

给定一个长度为$n$的数组$a$，求具有公差的子序列的方案数，子序列要求长度大于等于3。

Constraints

• $1\le n\le 1000$
• $-2^{31}\le a_i\le 2^{31}-1$

Solution

想到一个不好做的递推公式：设$dp[i][j]$，表示以$i$结尾且公差为$nums[i]-nums[j]$的子序列方案数。设想美好但是推不动，似乎没法在$dp[i-1][?]$反推$nums[j]$……

那改一改：设$dp[i][d]$，表示以$i$结尾的公差为$d$的子序列方案数。这下子第二维用哈希表维护就好了。

另外注意一下是求全部序列而不是只求$dp[n-1]$，还有必踩坑的溢出情况。

Code

class Solution {
public:
    int numberOfArithmeticSlices(vector<int>& nums) {
        // NOTE
        using L = long long;
        // dp:  len(seq) > 1
        // ans: len(seq) > 2
        unordered_map<L, int> dp[1003];
        int ans = 0;
        for(int i = 0; i < nums.size(); ++i) {
            for(int j = 0; j < i; ++j) {
                auto d = L(nums[i]) - nums[j];
                auto iter = dp[j].find(d);
                auto v = iter != dp[i-1].end() ?
                        iter->second : 0;
                // v: {any seq, nums[i]}
                // 1: {nums[j], nums[i]}
                dp[i][d] += v + 1;
                ans += v;
            }
        }
        return ans;
    }
};

Link

LeetCode-446 Arithmetic Slices II - Subsequence