We use cookies (including Google cookies) to personalize ads and analyze traffic. By continuing to use our site, you accept our Privacy Policy.

SimplyLeet
TagsCompaniesDisclaimersPrivacyTermsContactAbout Us

Count Non-Decreasing Subarrays After K Operations

Number: 3674

Difficulty: Hard

Paid? No

Companies: Microsoft, Google

Problem Description

Given an integer array nums and an integer k, each subarray can have at most k operations applied where one may increment any element by 1 per operation. For each subarray (considered independently), determine if one can transform it into a non-decreasing sequence by performing at most k operations. Return the count of subarrays for which this is possible.

A subarray is non-decreasing if for every adjacent pair the later element is greater than or equal to the previous element.

Key Insights

  • The minimum number of operations needed to “fix” a subarray is found by simulating the greedy process: start from the first element and ensure each subsequent element is at least the current maximum; if not, increment it. Thus, the cost for subarray nums[l...r] equals ∑₍ᵢ₌ₗ₊₁₎^r (max_{j=l..i-1} – nums[i]) (only when this difference is positive).
  • For any fixed starting index l, as we extend r the cost is non-decreasing (each new element only adds a non-negative cost).
  • If we let current = nums[l] and process the subarray, then when a new element is encountered: • If it is less than current, the cost increases by (current – nums[i]). • Otherwise, current is updated to nums[i] (and no extra cost is added).
  • One can “decompose” the subarray into segments where the maximum (thus the “target” value) is constant. In such a segment, the cost is computed as: segment_cost = (segment_length × value) – (sum of the values in that segment).
  • Precomputing the next-greater element for each index enables a segmentation of the subarray and helps compute cost quickly. In addition, constructing prefix sums allows us to calculate the cost over any segment in O(1) time.
  • Then, for every starting index l, we can use binary search (or a two‐pointer style “expanding window” that jumps over segments) to find the maximum valid r where the cost is ≤ k.
  • Finally, by summing over all l the number of valid subarrays that start at l, we obtain the answer.

Space and Time Complexity

Time Complexity: O(n log n) in the worst case. For each starting index, binary search is used to locate the maximum r. (A careful “two‐pointers” solution might achieve O(n) overall but requires sophisticated maintenance.) Space Complexity: O(n) for the prefix sums, next-greater indices array, and any auxiliary data structures.

Solution

We first note that for a subarray nums[l...r] the minimal cost to make it non‐decreasing is achieved by “fixing” it in one pass: starting with current = nums[l], for each index i from l+1 to r, if nums[i] is smaller than current then we would spend (current – nums[i]) operations and otherwise update current. This sum is our “cost” for that subarray.

One idea is to precompute for every index its “next greater index” (i.e. the first index to the right having a strictly larger value). This divides the subarray starting at l into segments where the maximum remains constant. Specifically, if for a starting index l the first segment runs up to index r₁ (where r₁ is the next greater position than l or the end of the array), then for every i in [l+1, r₁-1] the “target” value is nums[l]. The cost over this segment is:   nums[l] × (segment_length) – (sum of values in that segment).

If r extends into the next segment, the target is updated and the cost is computed similarly on the next segment. By precomputing prefix sums of nums the cost over any segment can be computed quickly.

For each starting index l, we then binary search over r (or move a pointer by segments) while accumulating the cost. The subarray is valid if the total cost does not exceed k. Finally, the answer is the total count of all valid subarrays.

The trick is in efficiently “querying” the cost for a given r from l. With the next-greater structure and prefix sums, we can compute the cost piecewise without scanning every element in the subarray.

A detailed implementation will:

  1. Precompute prefix sums.
  2. Precompute next greater elements for every index.
  3. For each starting index l, simulate the segmentation of the subarray by “jumping” segments and use binary search (or a while-loop with careful bounds) to determine the largest r such that cost(l, r) ≤ k.
  4. Sum over all l the number of valid subarrays starting at l.

Code Solutions

Below are implementations in Python, JavaScript, C++, and Java. Each solution includes line‐by‐line comments.

# Python solution

def countSubarrays(nums, k):
    n = len(nums)
    # Compute prefix sums: prefix[i] = sum of nums[0..i]
    prefix = [0] * (n + 1)
    for i in range(n):
        prefix[i+1] = prefix[i] + nums[i]
    
    # Precompute next greater element indices for each index.
    # next_greater[i] = smallest j > i such that nums[j] > nums[i]
    next_greater = [n] * n
    stack = []
    for i in range(n):
        while stack and nums[i] > nums[stack[-1]]:
            idx = stack.pop()
            next_greater[idx] = i
        stack.append(i)
    
    total = 0
    # For each starting index, find the maximum r where cost <= k.
    for l in range(n):
        cost = 0
        r = l
        current = nums[l]
        # Use a temporary pointer to move through segments.
        while r < n:
            # Determine the boundary of the current segment.
            seg_end = next_greater[r] if next_greater[r] < n else n
            # Now, in the segment starting at r, the current target value is current.
            # End index for this segment (within the subarray from l) is min(seg_end, n).
            # Binary search in the range [r, seg_end] to find the furthest index r2
            # such that cost added = current*(r2 - r + 1) - (prefix[r2+1] - prefix[r]) does not exceed k - cost.
            lo, hi = r, seg_end - 1
            valid = r - 1
            while lo <= hi:
                mid = (lo + hi) // 2
                added = current * (mid - r + 1) - (prefix[mid+1] - prefix[r])
                if cost + added <= k:
                    valid = mid
                    lo = mid + 1
                else:
                    hi = mid - 1
            # Count subarrays from starting index l ending at indices [r, valid]
            total += (valid - r + 1)
            # If valid reached the end of current segment then move to next segment if possible,
            # else break because we cannot extend further.
            if valid == seg_end - 1 and seg_end < n:
                # Include the element at seg_end: update cost and current.
                added = current * (seg_end - r) - (prefix[seg_end] - prefix[r])
                cost += added
                # Now, update current to new maximum.
                current = nums[seg_end]
                r = seg_end
            else:
                # Cannot extend further without exceeding the allowed operations.
                break
    return total

# Example usage:
nums = [6,3,1,2,4,4]
k = 7
print(countSubarrays(nums, k))  # Expected output: 17
← Back to All Questions