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