Scale Engineer

Count Partitions With Max-Min Difference at Most K

Med
#3186Time: 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.
Patterns
Sliding WindowDynamic ProgrammingPrefix Sum
Data structures
ArrayQueueMonotonic Queue

Prompt

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

3 approaches with complexity analysis and trade-offs.

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].

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].

Walkthrough

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.

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];    }}

Complexity

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.

Trade-offs

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.

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next →

Next problem