#3578
Medium Algorithms Count partitions with max min difference at most k
Array Dynamic Programming Queue Sliding Window Prefix Sum Monotonic Queue
58.8% acceptance
Feb 25, 2026
525
132
Partition nums into non-empty contiguous segments where max-min <= k in each segment.
Return the number of ways mod 10^9+7.
Solution
Rust
Time O(n²)
Space O(n)
impl Solution {
pub fn count_partitions(nums: Vec<i32>, k: i32) -> i32 {
const MOD: i64 = 1_000_000_007;
let n = nums.len();
let k = k as i64;
// dp[i] = number of ways to partition nums[0..i]
// dp[i] = sum of dp[j] for all j where max(j..i) - min(j..i) <= k
// Use sliding window with monotonic deques for max and min
let mut dp = vec![0i64; n + 1];
dp[0] = 1;
use std::collections::VecDeque;
let mut max_dq: VecDeque<usize> = VecDeque::new(); // decreasing for max
let mut min_dq: VecDeque<usize> = VecDeque::new(); // increasing for min
let mut window_sum = 0i64;
let mut left = 0usize; // leftmost valid start of current segment ending at i
// For each i (1-indexed end), the last segment is nums[j..=i-1] = nums[j..i]
// So j ranges from 0 to i-1 but the segment is [j..i-1] in 0-indexed nums.
// Let's think: dp[i] = sum dp[j] for j in [l..i] where nums[j..i-1] is valid (i is 1-indexed end).
for i in 1..=n {
// Add nums[i-1] to window
while !max_dq.is_empty() && nums[*max_dq.back().unwrap()] <= nums[i - 1] {
max_dq.pop_back();
}
max_dq.push_back(i - 1);
while !min_dq.is_empty() && nums[*min_dq.back().unwrap()] >= nums[i - 1] {
min_dq.pop_back();
}
min_dq.push_back(i - 1);
// Add dp[i-1] to window sum (j = i-1 means segment [i-1..i-1] = single element)
window_sum = (window_sum + dp[i - 1]) % MOD;
// Shrink window from left while max-min > k
while !max_dq.is_empty()
&& !min_dq.is_empty()
&& (nums[*max_dq.front().unwrap()] as i64 - nums[*min_dq.front().unwrap()] as i64) > k
{
// Remove dp[left] from window sum
window_sum = (window_sum - dp[left] + MOD) % MOD;
if *max_dq.front().unwrap() == left {
max_dq.pop_front();
}
if *min_dq.front().unwrap() == left {
min_dq.pop_front();
}
left += 1;
}
dp[i] = window_sum;
}
dp[n] as i32
}
}