3251. 单调数组对的数目 II - 力扣(LeetCode)

超时方案

class Solution {
public:
    int countOfPairs(vector<int>& nums) {
        const int MOD = 1000000007;
        int n = nums.size();
        int m = ranges::max(nums);
        vector dp(n + 1, vector<long long>(m + 1));
        for (int i = 0; i <= nums[0]; i++)
            dp[1][i] = 1;
        for (int i = 2; i <= n; i++) {
            for (int j = nums[i - 1]; j >= 0; j--) {
                for (int k = 0; nums[i - 2] - k >= nums[i - 1] - j && k <= j; k++) {
                    dp[i][j] = (dp[i][j] + dp[i - 1][k]) % MOD;
                }
            }
        }
        int ans = 0;
        for (int i = 0; i <= nums[n - 1]; i++)
            ans = (ans + dp[n][i]) % MOD;
        return ans;
    }
};

该方案可以通过单调数组对的数目 I,但是这一题数据范围更大,会超时

题解

class Solution {
public:
    int countOfPairs(vector<int>& nums) {
        const int MOD = 1000000007;
        int n = nums.size();
        int m = ranges::max(nums);
        vector dp(n + 1, vector<long long>(m + 1));
        vector<long long> s(m + 1); // 保存前一位置前缀和
        for (int i = 0; i <= nums[0]; i++) dp[1][i] = 1;
        for (int i = 2; i <= n; i++) {
            partial_sum(dp[i - 1].begin(), dp[i - 1].end(), s.begin()); // 计算前缀和 
            for (int j = nums[i - 1]; j >= 0; j--) { 
                int max_k = j + min(0, nums[i - 2] - nums[i - 1]);
                dp[i][j] = max_k >= 0 ? s[max_k] % MOD : 0; 
            }
        }
        int ans = 0;
        for (int i = 0; i <= nums[n - 1]; i++)
            ans = (ans + dp[n][i]) % MOD;
        return ans;
    }
};

对超时方案的优化,减少了每次都要计算前缀和的时间。

思路

DP+前缀和优化

DP[i][j]表示前i个nums中,最后一个是j时,数组对的个数
s是DP[i-1]的前缀和