LeetCode 3578 - Count Partitions With Max-Min Difference at Most K

We are given an array nums and an integer k. We want to split the array into one or more contiguous, non-empty segments. Every element must belong to exactly one segment, and the segments must preserve the original order of the array.

leetcodemediumarraydynamic-programmingqueuesliding-windowprefix-summonotonic-queue

LeetCode Problem 3578

Difficulty: 🟡 Medium
Topics: Array, Dynamic Programming, Queue, Sliding Window, Prefix Sum, Monotonic Queue

Solution

LeetCode 3578 - Count Partitions With Max-Min Difference at Most K

Problem Understanding

We are given an array nums and an integer k. We want to split the array into one or more contiguous, non-empty segments. Every element must belong to exactly one segment, and the segments must preserve the original order of the array.

A partition is considered valid if, for every segment, the difference between the largest element and the smallest element inside that segment is at most k.

Our goal is to count how many different valid partitions exist.

For example, if a segment contains values [4, 1, 3], then:

  • Maximum = 4
  • Minimum = 1
  • Difference = 3

This segment is valid if 3 <= k.

The output can become extremely large because the number of possible partitions grows exponentially with the array size. Therefore, we must return the answer modulo 10^9 + 7.

The constraints are especially important:

  • n = nums.length can be as large as 5 * 10^4
  • Values can be as large as 10^9
  • k can also be as large as 10^9

Since n reaches fifty thousand, any algorithm with quadratic complexity is too slow. We need a solution that is close to linear time.

Some important edge cases include:

  • k = 0, where every segment must contain identical values.
  • Arrays where every element can belong to one giant segment.
  • Arrays where no adjacent elements can coexist in the same segment.
  • Large inputs requiring careful modulo handling.
  • Repeated values that create many valid partition possibilities.

Approaches

Brute Force

The most direct approach is to try every possible partition.

For an array of length n, there are n - 1 positions between elements. At each position we either cut or do not cut, resulting in:

$$2^{n-1}$$

possible partitions.

For each partition, we can verify every segment by computing its minimum and maximum values.

This approach is correct because it explicitly checks every possible partition and counts the valid ones.

Unfortunately, it is completely infeasible. Even for n = 50, the number of partitions is already enormous, while the actual constraint is 50000.

Key Insight

Let:

  • dp[i] = number of valid ways to partition the prefix ending before index i.

More precisely:

  • dp[0] = 1
  • dp[i] counts ways to partition nums[0:i]

Suppose we want to compute dp[r+1].

Any valid last segment ending at index r must start at some position l.

Then:

$$dp[r+1] += dp[l]$$

for every valid starting position l.

The challenge becomes finding all valid segment starts efficiently.

A segment [l, r] is valid when:

$$\max(nums[l..r]) - \min(nums[l..r]) \le k$$

As r moves from left to right, we can maintain the smallest valid left boundary using a sliding window.

To efficiently maintain window maximum and minimum values, we use two monotonic deques:

  • decreasing deque for maxima
  • increasing deque for minima

Once we know the smallest valid left boundary left, every starting position in:

$$[left, r]$$

forms a valid ending segment.

Therefore:

$$dp[r+1] = \sum_{i=left}^{r} dp[i]$$

A prefix-sum array allows this range sum to be computed in O(1).

This transforms the problem into an O(n) solution.

Approach Comparison

Approach Time Complexity Space Complexity Notes
Brute Force O(2^n · n) O(n) Enumerates every partition
Optimal O(n) O(n) Sliding window + monotonic queues + DP + prefix sums

Algorithm Walkthrough

  1. Create a DP array where dp[i] represents the number of valid partitions for the prefix nums[0:i].
  2. Initialize dp[0] = 1. An empty prefix has exactly one way to be partitioned.
  3. Maintain two monotonic deques:
  • A decreasing deque storing indices of potential maximum values.
  • An increasing deque storing indices of potential minimum values.
  1. Maintain a sliding window [left, right].
  2. For each right from 0 to n-1:
  • Insert nums[right] into both monotonic deques.
  • Remove smaller elements from the maximum deque.
  • Remove larger elements from the minimum deque.
  1. Check whether the current window violates:

$$\max - \min > k$$

If it does, move left forward until the window becomes valid again. 7. While moving left, remove expired indices from the fronts of the deques. 8. After adjustment, every segment ending at right and starting anywhere in [left, right] is valid. 9. Therefore:

$$dp[right+1] = \sum_{i=left}^{right} dp[i]$$ 10. Maintain a prefix-sum array:

$$pref[i] = dp[0] + dp[1] + \dots + dp[i]$$

  1. Compute the range sum in O(1):

$$dp[right+1] = pref[right] - pref[left-1]$$

  1. Update the prefix sums modulo 10^9 + 7.
  2. After processing all positions, return dp[n].

Why it works

The sliding window invariant guarantees that [left, right] is the smallest valid window ending at right. Therefore every start position from left through right forms a valid final segment, and every earlier start position is invalid.

The DP recurrence counts all ways to partition the prefix before the chosen segment start. Since every valid partition ends with exactly one final segment, and every valid final segment contributes exactly once, the recurrence counts all valid partitions without duplication or omission.

Python Solution

from collections import deque
from typing import List

class Solution:
    def countPartitions(self, nums: List[int], k: int) -> int:
        MOD = 10**9 + 7
        n = len(nums)

        dp = [0] * (n + 1)
        prefix = [0] * (n + 1)

        dp[0] = 1
        prefix[0] = 1

        max_deque = deque()
        min_deque = deque()

        left = 0

        for right, value in enumerate(nums):
            while max_deque and nums[max_deque[-1]] <= value:
                max_deque.pop()
            max_deque.append(right)

            while min_deque and nums[min_deque[-1]] >= value:
                min_deque.pop()
            min_deque.append(right)

            while nums[max_deque[0]] - nums[min_deque[0]] > k:
                if max_deque[0] == left:
                    max_deque.popleft()
                if min_deque[0] == left:
                    min_deque.popleft()
                left += 1

            total = prefix[right]
            if left > 0:
                total -= prefix[left - 1]

            dp[right + 1] = total % MOD
            prefix[right + 1] = (prefix[right] + dp[right + 1]) % MOD

        return dp[n]

The implementation follows the algorithm directly.

The two deques maintain the maximum and minimum values inside the current sliding window. The left boundary only moves forward, which guarantees linear total deque operations.

For each position right, the valid starting positions for the last segment form a contiguous range [left, right]. Instead of summing dp[left] through dp[right] manually, the prefix-sum array computes the range sum in constant time.

The modulo is applied after each DP update and prefix update to prevent overflow and satisfy the problem requirements.

Go Solution

func countPartitions(nums []int, k int) int {
	const MOD int = 1000000007

	n := len(nums)

	dp := make([]int, n+1)
	prefix := make([]int, n+1)

	dp[0] = 1
	prefix[0] = 1

	maxDeque := []int{}
	minDeque := []int{}

	left := 0

	for right, value := range nums {
		for len(maxDeque) > 0 &&
			nums[maxDeque[len(maxDeque)-1]] <= value {
			maxDeque = maxDeque[:len(maxDeque)-1]
		}
		maxDeque = append(maxDeque, right)

		for len(minDeque) > 0 &&
			nums[minDeque[len(minDeque)-1]] >= value {
			minDeque = minDeque[:len(minDeque)-1]
		}
		minDeque = append(minDeque, right)

		for nums[maxDeque[0]]-nums[minDeque[0]] > k {
			if maxDeque[0] == left {
				maxDeque = maxDeque[1:]
			}
			if minDeque[0] == left {
				minDeque = minDeque[1:]
			}
			left++
		}

		total := prefix[right]
		if left > 0 {
			total -= prefix[left-1]
		}

		total %= MOD
		if total < 0 {
			total += MOD
		}

		dp[right+1] = total
		prefix[right+1] = (prefix[right] + dp[right+1]) % MOD
	}

	return dp[n]
}

The Go version mirrors the Python implementation closely. The deques are implemented using slices that store indices. Since subtraction can temporarily produce a negative value when computing the range sum, the result is normalized back into the modulo range before storing it in dp.

Worked Examples

Example 1

Input:

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

Iteration Details

right value left Valid starts dp[right+1]
0 9 0 [0] 1
1 4 0 [0,1] 2
2 1 1 [1,2] 3
3 3 1 [1,2,3] 6
4 7 4 [4] 6

DP progression:

i dp[i]
0 1
1 1
2 2
3 3
4 6
5 6

Answer:

6

Example 2

Input:

nums = [3,3,4]
k = 0

Iteration Details

right value left Valid starts dp[right+1]
0 3 0 [0] 1
1 3 0 [0,1] 2
2 4 2 [2] 2

DP progression:

i dp[i]
0 1
1 1
2 2
3 2

Answer:

2

Complexity Analysis

Measure Complexity Explanation
Time O(n) Each index enters and leaves each deque at most once
Space O(n) DP array, prefix array, and deques

The key observation is that the sliding window's left pointer only moves forward. Every element is pushed and popped from each monotonic deque at most one time, giving linear deque work overall. The DP transition becomes O(1) through prefix sums, yielding an overall O(n) algorithm.

Test Cases

from typing import List

s = Solution()

assert s.countPartitions([9,4,1,3,7], 4) == 6      # provided example 1
assert s.countPartitions([3,3,4], 0) == 2          # provided example 2

assert s.countPartitions([1,1], 0) == 2            # both partitions valid
assert s.countPartitions([1,10], 0) == 1           # must split

assert s.countPartitions([5,5,5,5], 0) == 8        # all cuts optional

assert s.countPartitions([1,2,3], 100) == 4        # every segment valid

assert s.countPartitions([1,100,1], 0) == 1        # every element isolated

assert s.countPartitions([1,2,1,2], 1) == 8        # every segment valid

assert s.countPartitions([1,3,6], 2) == 2          # only some merges valid

assert s.countPartitions([7,7,7,7,7], 0) == 16     # 2^(n-1)

assert s.countPartitions([1,5,9,13], 3) == 1       # every adjacent pair invalid

Test Summary

Test Why
[9,4,1,3,7], k=4 Official example
[3,3,4], k=0 Official example
[1,1], k=0 Smallest repeated-value case
[1,10], k=0 Smallest forced-split case
[5,5,5,5], k=0 Every cut position optional
[1,2,3], k=100 Entire array always valid
[1,100,1], k=0 Every element must stand alone
[1,2,1,2], k=1 Many overlapping valid segments
[1,3,6], k=2 Mixed valid and invalid windows
[7,7,7,7,7], k=0 Maximum partition flexibility
[1,5,9,13], k=3 No adjacent merge allowed

Edge Cases

k Equals Zero

When k = 0, every valid segment must consist entirely of equal values. This can easily cause bugs if the window maintenance logic incorrectly allows a larger range. The monotonic deques accurately track the current maximum and minimum values, so the condition max - min <= 0 is enforced exactly.

Every Segment Is Valid

If k is extremely large, the entire array may always satisfy the constraint. In this situation the left pointer never moves, and every previous partition state contributes to the current one. The prefix-sum recurrence naturally accumulates all possibilities and effectively computes the expected 2^(n-1) count.

Every Element Must Be Isolated

Consider arrays where any pair of neighboring elements already violates the constraint. Then the window shrinks immediately whenever a second element is added, causing left = right. The DP transition becomes dp[right+1] = dp[right], meaning there is exactly one valid partition. The implementation handles this automatically through the sliding-window adjustment process.

Large Input Size

With n = 50000, quadratic solutions are impossible. The implementation ensures that each index is inserted and removed from each deque at most once. Combined with O(1) DP transitions via prefix sums, the algorithm remains efficient even at the maximum constraint size.

Large Answer Values

The number of valid partitions can grow exponentially. Without modular arithmetic, integer values would become enormous. Every DP update and prefix-sum update is performed modulo 10^9 + 7, guaranteeing correctness and preventing overflow.