Count Partitions With Max-Min Difference at Most K

medium monotonic queue sliding window dynamic programming prefix sum

Problem

Split nums into one or more non-empty contiguous segments so that in every segment max − min ≤ k. Count the number of valid partitions, returned modulo 10⁹ + 7.

Let dp[i] be the number of valid partitions of the prefix of length i. A new segment can end at position i and start at any j for which the window nums[j..i−1] still satisfies max − min ≤ k. Those starts form a contiguous range whose left edge only moves right, so a sliding window with two monotonic deques (one tracking the window max, one the min) plus a prefix sum of dp gives each dp[i] in amortized O(1).

Inputnums = [9, 4, 1, 3, 7], k = 4

Output6

Six partitions keep every segment within a spread of 4, e.g. [[9],[4,1],[3,7]] and [[9],[4,1,3],[7]].

nums (comma-separated)

def countPartitions(nums, k):
    MOD = 10**9 + 7
    n = len(nums)
    dp = [0] * (n + 1)          # dp[i] = ways for prefix length i
    dp[0] = 1
    pre = [0] * (n + 1)         # pre[i] = dp[0] + ... + dp[i]
    pre[0] = 1
    maxq, minq = deque(), deque()   # decreasing / increasing indices
    left = 0                    # earliest valid segment start
    for r in range(n):
        while maxq and nums[maxq[-1]] <= nums[r]:
            maxq.pop()
        maxq.append(r)
        while minq and nums[minq[-1]] >= nums[r]:
            minq.pop()
        minq.append(r)
        while nums[maxq[0]] - nums[minq[0]] > k:   # shrink window
            if maxq[0] < minq[0]:
                left = maxq.popleft() + 1
            else:
                left = minq.popleft() + 1
        # dp[r+1] = sum of dp[left .. r]
        s = pre[r] - (pre[left - 1] if left > 0 else 0)
        dp[r + 1] = s % MOD
        pre[r + 1] = (pre[r] + dp[r + 1]) % MOD
    return dp[n] % MOD

function countPartitions(nums, k) {
  const MOD = 1000000007n, n = nums.length;
  const dp = new Array(n + 1).fill(0n);   // dp[i] = ways for prefix i
  dp[0] = 1n;
  const pre = new Array(n + 1).fill(0n);  // pre[i] = dp[0..i] sum
  pre[0] = 1n;
  const maxq = [], minq = [];             // decreasing / increasing
  let left = 0;                           // earliest valid start
  for (let r = 0; r < n; r++) {
    while (maxq.length && nums[maxq[maxq.length-1]] <= nums[r]) maxq.pop();
    maxq.push(r);
    while (minq.length && nums[minq[minq.length-1]] >= nums[r]) minq.pop();
    minq.push(r);
    while (nums[maxq[0]] - nums[minq[0]] > k) {   // shrink window
      if (maxq[0] < minq[0]) left = maxq.shift() + 1;
      else left = minq.shift() + 1;
    }
    // dp[r+1] = sum of dp[left .. r]
    const s = pre[r] - (left > 0 ? pre[left - 1] : 0n);
    dp[r + 1] = ((s % MOD) + MOD) % MOD;
    pre[r + 1] = (pre[r] + dp[r + 1]) % MOD;
  }
  return Number(dp[n] % MOD);
}

int countPartitions(int[] nums, int k) {
    long MOD = 1_000_000_007L;
    int n = nums.length;
    long[] dp = new long[n + 1];        // dp[i] = ways for prefix i
    dp[0] = 1;
    long[] pre = new long[n + 1];       // pre[i] = dp[0..i] sum
    pre[0] = 1;
    Deque<Integer> maxq = new ArrayDeque<>();  // decreasing
    Deque<Integer> minq = new ArrayDeque<>();  // increasing
    int left = 0;                       // earliest valid start
    for (int r = 0; r < n; r++) {
        while (!maxq.isEmpty() && nums[maxq.peekLast()] <= nums[r]) maxq.pollLast();
        maxq.addLast(r);
        while (!minq.isEmpty() && nums[minq.peekLast()] >= nums[r]) minq.pollLast();
        minq.addLast(r);
        while ((long) nums[maxq.peekFirst()] - nums[minq.peekFirst()] > k) {
            if (maxq.peekFirst() < minq.peekFirst()) left = maxq.pollFirst() + 1;
            else left = minq.pollFirst() + 1;
        }
        long s = pre[r] - (left > 0 ? pre[left - 1] : 0);
        dp[r + 1] = ((s % MOD) + MOD) % MOD;
        pre[r + 1] = (pre[r] + dp[r + 1]) % MOD;
    }
    return (int) (dp[n] % MOD);
}

int countPartitions(vector<int>& nums, int k) {
    const long long MOD = 1000000007LL;
    int n = nums.size();
    vector<long long> dp(n + 1, 0), pre(n + 1, 0);
    dp[0] = 1; pre[0] = 1;              // dp[i], prefix sums
    deque<int> maxq, minq;             // decreasing / increasing
    int left = 0;                       // earliest valid start
    for (int r = 0; r < n; r++) {
        while (!maxq.empty() && nums[maxq.back()] <= nums[r]) maxq.pop_back();
        maxq.push_back(r);
        while (!minq.empty() && nums[minq.back()] >= nums[r]) minq.pop_back();
        minq.push_back(r);
        while ((long long) nums[maxq.front()] - nums[minq.front()] > k) {
            if (maxq.front() < minq.front()) left = maxq.front() + 1, maxq.pop_front();
            else left = minq.front() + 1, minq.pop_front();
        }
        long long s = pre[r] - (left > 0 ? pre[left - 1] : 0);
        dp[r + 1] = ((s % MOD) + MOD) % MOD;
        pre[r + 1] = (pre[r] + dp[r + 1]) % MOD;
    }
    return (int) (dp[n] % MOD);
}

Explanation

We count partitions with a left-to-right dynamic program. Let dp[i] be the number of valid ways to partition the first i elements. The answer is dp[n], and dp[0] = 1 (the empty prefix has exactly one partition: nothing).

To build dp[i] we decide where the last segment begins. If that segment is nums[j..i−1], it is allowed only when its max − min ≤ k, and the rest contributes dp[j] ways. So dp[i] = ∑ dp[j] over every legal start j.

The crucial observation is that for a fixed right end r = i−1, the legal starts form a contiguous range [left, r]: extending a window can only break the spread condition, never repair it, so as r grows left only moves right. That is a classic sliding window.

To know the window's max and min in O(1) we keep two monotonic deques. maxq stores indices with decreasing values (its front is the window maximum); minq stores indices with increasing values (its front is the window minimum). When the spread exceeds k we drop whichever front is older and push left past it.

Finally, dp[i] is a range sum dp[left] + … + dp[r]. A running prefix-sum array pre, where pre[i] = dp[0] + … + dp[i], turns that into a single subtraction. Each index enters and leaves each deque at most once, so the whole scan is linear.

Example: nums = [9,4,1,3,7], k = 4. The windows ending at each position allow exactly the starts that keep the spread within 4, and summing the corresponding dp values yields dp[5] = 6.

Time: O(n) Space: O(n)