# Count Partitions With Max-Min Difference at Most K
**Difficulty:** MEDIUM
[External](https://leetcode.com/problems/count-partitions-with-max-min-difference-at-most-k)
Canonical: https://scaleengineer.com/dsa/problems/count-partitions-with-max-min-difference-at-most-k
**Patterns:** [Sliding Window](https://scaleengineer.com/dsa/patterns/sliding-window), [Dynamic Programming](https://scaleengineer.com/dsa/patterns/dynamic-programming), [Prefix Sum](https://scaleengineer.com/dsa/patterns/prefix-sum)
**Data structures:** Array, Queue, Monotonic Queue
---
## Problem
You are given an integer array `nums` and an integer `k`. Your task is to partition `nums` into one or more **non-empty** contiguous segments such that in each segment, the difference between its **maximum** and **minimum** elements is **at most** `k`.

Return the total number of ways to partition `nums` under this condition.

Since the answer may be too large, return it **modulo** `109 + 7`.

**Example 1:**

**Input:** nums = \[9,4,1,3,7\], k = 4

**Output:** 6

**Explanation:**

There are 6 valid partitions where the difference between the maximum and minimum elements in each segment is at most `k = 4`:

* `[[9], [4], [1], [3], [7]]`
* `[[9], [4], [1], [3, 7]]`
* `[[9], [4], [1, 3], [7]]`
* `[[9], [4, 1], [3], [7]]`
* `[[9], [4, 1], [3, 7]]`
* `[[9], [4, 1, 3], [7]]`

**Example 2:**

**Input:** nums = \[3,3,4\], k = 0

**Output:** 2

**Explanation:**

There are 2 valid partitions that satisfy the given conditions:

* `[[3], [3], [4]]`
* `[[3, 3], [4]]`

**Constraints:**

* `2 <= nums.length <= 5 * 104`
* `1 <= nums[i] <= 109`
* `0 <= k <= 109`

# Approaches
## Dynamic Programming
This approach uses dynamic programming to solve the problem. We define `dp[i]` as the number of ways to partition the prefix of the array `nums[0...i-1]`. To compute `dp[i]`, we consider all possible last segments ending at index `i-1`. If a segment `nums[j...i-1]` is valid (i.e., its max-min difference is at most `k`), we can form a valid partition by appending this segment to any valid partition of `nums[0...j-1]`. The number of ways to do this is `dp[j]`. By summing up `dp[j]` for all valid `j`, we get `dp[i]`.
**Time:** O(n^2) because of the nested loops. The outer loop runs `n` times, and the inner loop can run up to `n` times in the worst case. · **Space:** O(n) to store the `dp` array.
**Pros:** * It's a straightforward implementation of the DP recurrence.; * The logic is relatively easy to understand.
**Cons:** * The time complexity of `O(n^2)` is too slow for the given constraints (`n <= 5 * 10^4`), and will likely result in a Time Limit Exceeded (TLE) error.
### Explanation
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.

```java
class Solution {
    public int countPartitions(int[] nums, int k) {
        int n = nums.length;
        int MOD = 1_000_000_007;
        long[] dp = new long[n + 1];
        dp[0] = 1;

        for (int i = 1; i <= n; i++) {
            long currentMax = nums[i-1];
            long currentMin = nums[i-1];
            // j is the start index of the last segment
            for (int j = i - 1; j >= 0; j--) {
                currentMax = Math.max(currentMax, nums[j]);
                currentMin = Math.min(currentMin, nums[j]);
                if (currentMax - currentMin <= k) {
                    // Last segment is nums[j...i-1]
                    // Number of ways is dp[j]
                    dp[i] = (dp[i] + dp[j]) % MOD;
                } else {
                    // Since max-min is monotonic as j decreases,
                    // we can break early.
                    break;
                }
            }
        }
        return (int) dp[n];
    }
}
```
### Algorithm
- 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]`.

## DP with Binary Search and Segment Tree
This approach optimizes the `O(n^2)` DP solution. We observe that for a fixed end position `i-1`, the valid start positions `j` form a contiguous block `[p, i-1]`. This is because if `nums[p...i-1]` is valid, any subsegment `nums[j...i-1]` (where `j > p`) will also be valid. This allows us to rewrite the recurrence as `dp[i] = dp[p] + ... + dp[i-1]`. This sum can be calculated efficiently using prefix sums on the `dp` array.

The main challenge is to find the smallest valid start index `p` for each `i`. Since the validity condition is monotonic with respect to the start index `j`, we can use binary search to find `p`. To make the check inside the binary search efficient (i.e., finding max/min in a range), we use a data structure like a Segment Tree, which can answer range queries in logarithmic time.
**Time:** O(n log^2 n). The main loop runs `n` times. Inside, binary search takes `O(log n)` steps, and each step involves a segment tree query of `O(log n)`. · **Space:** O(n) for the `dp` array, prefix sum array, and the segment tree.
**Pros:** * Much faster than the naive `O(n^2)` DP.; * Demonstrates the use of powerful data structures to optimize DP transitions.
**Cons:** * The implementation is more complex, requiring both a segment tree and binary search within the DP loop.; * While faster than `O(n^2)`, it may still be too slow for the tightest time limits.
### Explanation
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`.

```java
class Solution {
    // Assume a SegmentTree class is implemented with a method:
    // query(left, right) that returns {min, max} in O(log n).
    SegmentTree st;

    public int countPartitions(int[] nums, int k) {
        int n = nums.length;
        st = new SegmentTree(nums); // O(n) build
        int MOD = 1_000_000_007;

        long[] dp = new long[n + 1];
        long[] ps = new long[n + 2];
        dp[0] = 1;
        ps[1] = 1;

        for (int i = 1; i <= n; i++) {
            // Binary search for the smallest valid starting index p
            int low = 0, high = i - 1;
            int p = i; 
            while (low <= high) {
                int mid = low + (high - low) / 2;
                long[] range = st.query(mid, i - 1); // {min, max}
                if (range[1] - range[0] <= k) {
                    p = mid;
                    high = mid - 1;
                } else {
                    low = mid + 1;
                }
            }
            
            if (p < i) {
                long ways = (ps[i] - ps[p] + MOD) % MOD;
                dp[i] = ways;
            }
            
            ps[i + 1] = (ps[i] + dp[i]) % MOD;
        }
        return (int) dp[n];
    }
}
```
### Algorithm
- 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)`.

## DP with Sliding Window and Monotonic Deques
This is the most efficient approach, achieving a linear time complexity. It builds upon the same DP formulation but finds the valid range of starting positions in a much faster way. Instead of re-computing or searching for the boundary `p` for each `i` from scratch, we use a sliding window. The right end of the window is the current index `i-1`, and the left end `p` is adjusted as we iterate. To find the max and min in this sliding window in amortized constant time, we use a classic technique involving two monotonic deques. One deque tracks the maximum in the current window, and the other tracks the minimum. This allows us to find `p` and compute `dp[i]` in a single pass through the array.
**Time:** O(n). Each index is pushed onto and popped from each deque at most once. The `left` pointer only moves forward. Thus, the work inside the loop is amortized O(1). · **Space:** O(n) for the `dp` array, prefix sum array, and the deques. In the worst case, the deques can hold up to `n` elements.
**Pros:** * Optimal `O(n)` time complexity, which is very efficient and passes all constraints.; * Uses a clever combination of sliding window and monotonic deques.
**Cons:** * The logic for maintaining the sliding window and monotonic deques can be complex and tricky to implement correctly.
### Explanation
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.

```java
import java.util.ArrayDeque;
import java.util.Deque;

class Solution {
    public int countPartitions(int[] nums, int k) {
        int n = nums.length;
        int MOD = 1_000_000_007;

        long[] dp = new long[n + 1];
        long[] ps = new long[n + 2];
        dp[0] = 1;
        ps[1] = 1;

        Deque<Integer> maxDeque = new ArrayDeque<>();
        Deque<Integer> minDeque = new ArrayDeque<>();
        int left = 0;

        for (int i = 1; i <= n; i++) {
            int currentIndex = i - 1;

            while (!maxDeque.isEmpty() && nums[maxDeque.peekLast()] <= nums[currentIndex]) {
                maxDeque.pollLast();
            }
            maxDeque.addLast(currentIndex);

            while (!minDeque.isEmpty() && nums[minDeque.peekLast()] >= nums[currentIndex]) {
                minDeque.pollLast();
            }
            minDeque.addLast(currentIndex);

            while (nums[maxDeque.peekFirst()] - nums[minDeque.peekFirst()] > k) {
                left++;
                if (maxDeque.peekFirst() < left) {
                    maxDeque.pollFirst();
                }
                if (minDeque.peekFirst() < left) {
                    minDeque.pollFirst();
                }
            }

            long ways = (ps[i] - ps[left] + MOD) % MOD;
            dp[i] = ways;
            
            ps[i + 1] = (ps[i] + dp[i]) % MOD;
        }

        return (int) dp[n];
    }
}
```
### Algorithm
- 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.