3578. 统计極差最大為 K 的分割方式數

思路

題目要把數字切成多段，每段的最大值 - 最小值 < k，問有幾種切法。
定義dp[i]代表0 ~ i這段共有幾種切法，可以寫出以下程式碼，

超時 dp

class Solution {
public:
    int countPartitions(vector<int>& nums, int k) {
        int n = nums.size();
        const int MOD = 1e9 + 7;
        vector<int> dp(n + 1);
        dp[0] = 1;
        for(int i = 1; i <= n; ++i) {
            int curMin = nums[i - 1], curMax = nums[i - 1];
            for(int j = i - 1; j >= 0; --j) { 
                curMin = min(curMin, nums[j]);
                curMax = max(curMax, nums[j]); // nums[j]  ~ nums[i] 的最大最小值。
                if(curMax - curMin <= k) { // 相減 <= k, 滿足條件，從 nums[j] ~ nums[i] 可以自成一段。
                    dp[i] += dp[j];
                    dp[i] %= MOD;
                }
                else { // 相減 > k, curMax - curMin 只會遞增，因為curMax只會越來越大，curMin 只會越來越小, 再往左擴已經沒有希望了。
                    break;
                }
            }
        }
        return dp[n];
    }
};

優化

可以注意到，當往左擴張時

• curMax 單調遞增
• curMin 單調遞減
• (curMax - curMin) 單調遞增 因此有加入break的操作。

往左擴張時，都會是連續滿足條件，到一個位置不滿足，其中的所有dp[j]都會被加到dp[i]上面。
因此優化能從這個地方開始，只要找到滿足條件，最左邊的位置left，然後將這段的dp數值總和，就找到答案了。
現在的問題是，要怎麼快速找到左端的位置。

我們能用滑動窗口去定位左端（天才！），窗口內部的dp總和就是答案，
為了要及時知道最大最小值，並且能及時找到第二大，第三大的數值，因此使用deque去紀錄。

由於需要計算窗口內部的dp總和，這會是一段連續的數值加總，因此能用前綴和去計算總合。

又因為每次計算sumDP[i]時，只會使用到sumDP[i - 1]的內容，可以用單一變數表達，以省去空間。

程式碼

1. 動態規劃 + 滑動窗口 + 雙端隊列 + 前綴和

class Solution {
public:
    int countPartitions(vector<int>& nums, int k) {
        int n = nums.size();
        const int MOD = 1e9 + 7;
        
        vector<long long> dp(n + 1), sumDP(n + 1);
        dp[0] = 1;
        sumDP[0] = 1;
        
        deque<int> max_q; // 對應到的數值單調遞減, 開頭是最大值的位置
        deque<int> min_q; // 對應到的數值單調遞增, 開頭是最小值的位置
        
        int left = 0;
        
        // max_q.push_back(0), min_q.push_back(0);
        for(int i = 1; i <= n; ++i) {            
            // nums[i - 1] 進入
            // 前面放最大值max_queue
            while(!max_q.empty() && nums[max_q.back()] <= nums[i - 1]) {
                max_q.pop_back();
            }
            max_q.push_back(i - 1);
            
            // 前面放最小值
            while(!min_q.empty() && nums[min_q.back()] >= nums[i - 1]) {
                min_q.pop_back();
            }
            min_q.push_back(i - 1);

            // 不滿足條件
            while(nums[max_q.front()] - nums[min_q.front()] > k) {
                if(max_q.front() == left) {
                    max_q.pop_front();
                }
                if(min_q.front() == left) {
                    min_q.pop_front();
                }
                ++left;
            }

            long long prev_sum = sumDP[i - 1];
            long long start_sum = (left > 0) ? sumDP[left - 1] : 0;
            dp[i] = (prev_sum - start_sum) % MOD;

            sumDP[i] = sumDP[i - 1] + dp[i];
        }
        return dp[n];
    }
};

2. 前綴和優化

class Solution {
public:
    int countPartitions(vector<int>& nums, int k) {
        const int MOD = 1e9 + 7;
        int n = nums.size();

        deque<int> min_q, max_q;
        vector<int> dp(n + 1);
        dp[0] = 1;
        
        long long sumDP = 0;
        int left = 0;
        for(int i = 0; i < n; ++i) {
            int x = nums[i];
            sumDP += dp[i];

            while(!min_q.empty() && nums[min_q.back()] >= x) {
                min_q.pop_back();
            }
            min_q.push_back(i);

            while(!max_q.empty() && nums[max_q.back()] <= x) {
                max_q.pop_back();
            }
            max_q.push_back(i);

            while(nums[max_q.front()] - nums[min_q.front()] > k) {
                sumDP -= dp[left];
                ++left;
                if(min_q.front() < left) {
                    min_q.pop_front();
                }
                if(max_q.front() < left) {
                    max_q.pop_front();
                }
            }
            dp[i + 1] = sumDP % MOD;
        }
        return dp[n];
    }
};

複雜度分析

• 時間複雜度：$O(n)$
• 空間複雜度：$O(n)$