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;
    }
};

LeetCode-446 Arithmetic Slices II - Subsequence