Count Partitions With Max-Min Difference at Most K

• Let dp[i] be the number of ways to partition the prefix nums[0...i-1].
• The base case is dp[0] = 1, representing one way to partition an empty prefix (by doing nothing).
• The recurrence relation is dp[i] = sum(dp[j]) for all 0 <= j < i such that the segment nums[j...i-1] is valid.
• A segment nums[j...i-1] is valid if max(nums[j...i-1]) - min(nums[j...i-1]) <= k.
• We can compute dp[i] for i from 1 to n:
  • For each i, iterate j from i-1 down to 0.
  • In this inner loop, maintain the maximum and minimum values of the current segment nums[j...i-1].
  • If the segment is valid, add dp[j] to dp[i] (modulo 10^9 + 7).
  • If the segment becomes invalid, we can break the inner loop because any segment starting before j will also be invalid (as the max-min difference is non-decreasing as the segment grows).
• The final answer is dp[n].

We build a dp array of size n+1, where dp[i] stores the number of valid partitions for the first i elements of nums. We initialize dp[0] = 1.

We then iterate from i = 1 to n. For each i, we need to calculate dp[i]. We do this by iterating backwards from j = i-1 to 0. The subarray nums[j...i-1] represents the last segment of a potential partition of nums[0...i-1].

In the inner loop, we keep track of the maximum and minimum elements in the segment nums[j...i-1]. If max - min <= k, it means this is a valid last segment. The number of ways to partition the rest of the array nums[0...j-1] is dp[j], so we add dp[j] to dp[i]. If max - min > k, we can stop the inner loop for the current i. This is because as j decreases, the segment nums[j...i-1] expands, and the difference between the maximum and minimum can only increase or stay the same. Therefore, all segments starting before j will also be invalid.

1class Solution {2    public int countPartitions(int[] nums, int k) {3        int n = nums.length;4        int MOD = 1_000_000_007;5        long[] dp = new long[n + 1];6        dp[0] = 1;7 8        for (int i = 1; i <= n; i++) {9            long currentMax = nums[i-1];10            long currentMin = nums[i-1];11            // j is the start index of the last segment12            for (int j = i - 1; j >= 0; j--) {13                currentMax = Math.max(currentMax, nums[j]);14                currentMin = Math.min(currentMin, nums[j]);15                if (currentMax - currentMin <= k) {16                    // Last segment is nums[j...i-1]17                    // Number of ways is dp[j]18                    dp[i] = (dp[i] + dp[j]) % MOD;19                } else {20                    // Since max-min is monotonic as j decreases,21                    // we can break early.22                    break;23                }24            }25        }26        return (int) dp[n];27    }28}

• The core DP recurrence is dp[i] = sum(dp[j]) for a range of valid start indices j.
• The set of valid start indices j for a segment ending at i-1 forms a contiguous range, say [p, i-1].
• The problem reduces to finding the smallest valid start index p for each i.
• We can compute the sum dp[p] + ... + dp[i-1] in O(1) using a prefix sum array over dp.
• To find p, we can use binary search over the possible start indices [0, i-1].
• For each mid in the binary search, we need to check if the segment nums[mid...i-1] is valid. This check involves finding the max and min in that range.
• We can pre-build a Segment Tree on nums to answer range max/min queries in O(log n) time.
• The binary search takes O(log n) steps. Each step involves an O(log n) query. So, finding p for each i takes O(log^2 n).
• The total time complexity is O(n * log^2 n).

First, we build a Segment Tree on the input array nums. This allows us to query the maximum and minimum values in any range [l, r] in O(log n) time. The build process takes O(n) time.

We maintain a dp array and a prefix sum array ps over dp. ps[i] will store dp[0] + ... + dp[i-1].

We iterate i from 1 to n. For each i, we perform a binary search on indices j from 0 to i-1 to find the smallest p such that max(nums[p...i-1]) - min(nums[p...i-1]) <= k.

• The binary search works as follows: for a mid index, we query our segment tree for the max and min in nums[mid...i-1].
• If max - min <= k, it means mid is a valid starting point, and there might be an even smaller valid index, so we search in the left half [low, mid-1] and record mid as a potential answer.
• If max - min > k, mid is not a valid start, so we must search in the right half [mid+1, high].

Once we find p, we calculate dp[i] = (ps[i] - ps[p] + MOD) % MOD. Then we update the prefix sum array: ps[i+1] = (ps[i] + dp[i]) % MOD.

1class Solution {2    // Assume a SegmentTree class is implemented with a method:3    // query(left, right) that returns {min, max} in O(log n).4    SegmentTree st;5 6    public int countPartitions(int[] nums, int k) {7        int n = nums.length;8        st = new SegmentTree(nums); // O(n) build9        int MOD = 1_000_000_007;10 11        long[] dp = new long[n + 1];12        long[] ps = new long[n + 2];13        dp[0] = 1;14        ps[1] = 1;15 16        for (int i = 1; i <= n; i++) {17            // Binary search for the smallest valid starting index p18            int low = 0, high = i - 1;19            int p = i; 20            while (low <= high) {21                int mid = low + (high - low) / 2;22                long[] range = st.query(mid, i - 1); // {min, max}23                if (range[1] - range[0] <= k) {24                    p = mid;25                    high = mid - 1;26                } else {27                    low = mid + 1;28                }29            }30            31            if (p < i) {32                long ways = (ps[i] - ps[p] + MOD) % MOD;33                dp[i] = ways;34            }35            36            ps[i + 1] = (ps[i] + dp[i]) % MOD;37        }38        return (int) dp[n];39    }40}

• This approach also uses the DP relation dp[i] = sum_{j=p}^{i-1} dp[j] and prefix sums.
• The key is to find the smallest valid start index p for each i in amortized O(1) time.
• We use a sliding window approach. We iterate i from 1 to n, which acts as the right end of our window. We maintain a left pointer, which corresponds to p.
• As i increments, left only moves to the right or stays put.
• To get the max and min of the sliding window [left, i-1] efficiently, we use two monotonic deques (double-ended queues).
  • maxDeque: Stores indices in decreasing order of their nums values.
  • minDeque: Stores indices in increasing order of their nums values.
• For each i:
  1. Update the deques with the new element nums[i-1].
  2. Check the window's validity: nums[maxDeque.front()] - nums[minDeque.front()].
  3. If the window is invalid (> k), shrink it from the left by incrementing left and removing indices from the deques that are no longer in the window.
  4. Once the window is valid, left is our p. Calculate dp[i] using prefix sums.
• The total time is linear because each element is added to and removed from the deques at most once.

We use the same DP state dp[i] and prefix sum array ps as in the previous approach. We also initialize a left pointer to 0, which will track the start of our valid window.

We iterate i from 1 to n. In each iteration, i-1 is the index of the new element being added to the right of our window.

1. Update Deques: We add i-1 to our two deques. For the maxDeque, we remove all indices from the end that correspond to values smaller than or equal to nums[i-1]. For the minDeque, we remove indices corresponding to values larger than or equal to nums[i-1]. This maintains the monotonic property of the deques.

2. Shrink Window: The maximum in the current window [left, i-1] is nums[maxDeque.peekFirst()] and the minimum is nums[minDeque.peekFirst()]. We check if their difference exceeds k. If it does, the window is invalid, so we must shrink it from the left by incrementing left. When we increment left, we also check if the indices at the front of the deques have fallen out of the new window boundary (< left) and remove them if so.

3. Calculate DP: After the shrinking step, left is the smallest index p such that nums[p...i-1] is a valid segment. We can then calculate dp[i] = (ps[i] - ps[left] + MOD) % MOD.

4. Update Prefix Sum: Finally, we update ps[i+1] = (ps[i] + dp[i]) % MOD.

This process ensures that each element is processed in amortized constant time, leading to an overall linear time solution.

1import java.util.ArrayDeque;2import java.util.Deque;3 4class Solution {5    public int countPartitions(int[] nums, int k) {6        int n = nums.length;7        int MOD = 1_000_000_007;8 9        long[] dp = new long[n + 1];10        long[] ps = new long[n + 2];11        dp[0] = 1;12        ps[1] = 1;13 14        Deque<Integer> maxDeque = new ArrayDeque<>();15        Deque<Integer> minDeque = new ArrayDeque<>();16        int left = 0;17 18        for (int i = 1; i <= n; i++) {19            int currentIndex = i - 1;20 21            while (!maxDeque.isEmpty() && nums[maxDeque.peekLast()] <= nums[currentIndex]) {22                maxDeque.pollLast();23            }24            maxDeque.addLast(currentIndex);25 26            while (!minDeque.isEmpty() && nums[minDeque.peekLast()] >= nums[currentIndex]) {27                minDeque.pollLast();28            }29            minDeque.addLast(currentIndex);30 31            while (nums[maxDeque.peekFirst()] - nums[minDeque.peekFirst()] > k) {32                left++;33                if (maxDeque.peekFirst() < left) {34                    maxDeque.pollFirst();35                }36                if (minDeque.peekFirst() < left) {37                    minDeque.pollFirst();38                }39            }40 41            long ways = (ps[i] - ps[left] + MOD) % MOD;42            dp[i] = ways;43            44            ps[i + 1] = (ps[i] + dp[i]) % MOD;45        }46 47        return (int) dp[n];48    }49}