3420. Count Non-Decreasing Subarrays After K Operations

Approach 1: Store (num, count)

• Time: $O(n)$
• Space: $O(n)$

C++JavaPython

class Solution {
 public:
  long long countNonDecreasingSubarrays(vector<int>& nums, int k) {
    long ans = 0;
    long cost = 0;
    // Store (number, count) pairs in non-increasing order. The numbers in the
    // queue represent what nums[i..j] look like after adjustments.
    deque<pair<int, int>> dq;

    for (int i = nums.size() - 1, j = nums.size() - 1; i >= 0; --i) {
      const int num = nums[i];
      int count = 1;
      while (!dq.empty() && dq.back().first < num) {
        const auto [nextNum, nextCount] = dq.back();
        dq.pop_back();
        count += nextCount;
        // Adjust `nextNum`s to `num`.
        cost += static_cast<long>(num - nextNum) * nextCount;
      }
      dq.emplace_back(num, count);
      while (cost > k) {
        const auto [rightmostNum, rightmostCount] = dq.front();
        dq.pop_front();
        cost -= static_cast<long>(rightmostNum - nums[j--]);
        if (rightmostCount > 1)
          dq.emplace_front(rightmostNum, rightmostCount - 1);
      }
      ans += j - i + 1;
    }

    return ans;
  }
};

class Solution {
  public long countNonDecreasingSubarrays(int[] nums, int k) {
    long ans = 0;
    long cost = 0;
    // Store (number, count) pairs in non-increasing order. The numbers in the
    // queue represent what nums[i..j] look like after adjustments.
    Deque<Pair<Integer, Integer>> dq = new ArrayDeque<>();

    for (int i = nums.length - 1, j = nums.length - 1; i >= 0; --i) {
      final int num = nums[i];
      int count = 1;
      while (!dq.isEmpty() && dq.getLast().getKey() < num) {
        final int nextNum = dq.getLast().getKey();
        final int nextCount = dq.removeLast().getValue();
        count += nextCount;
        // Adjust `nextNum`s to `num`.
        cost += (long) (num - nextNum) * nextCount;
      }
      dq.offerLast(new Pair<>(num, count));
      while (cost > k) {
        final int rightmostNum = dq.getFirst().getKey();
        final int rightmostCount = dq.removeFirst().getValue();
        cost -= (long) (rightmostNum - nums[j--]);
        if (rightmostCount > 1)
          dq.offerFirst(new Pair<>(rightmostNum, rightmostCount - 1));
      }
      ans += j - i + 1;
    }

    return ans;
  }
}

class Solution:
  def countNonDecreasingSubarrays(self, nums: list[int], k: int) -> int:
    ans = 0
    cost = 0
    # Store (number, count) pairs in non-increasing order. The numbers in the
    # queue represent what nums[i..j] look like after adjustments.
    dq = collections.deque()

    j = len(nums) - 1
    for i, num in reversed(list(enumerate(nums))):
      count = 1
      while dq and dq[-1][0] < num:
        nextNum, nextCount = dq.pop()
        count += nextCount
        cost += (num - nextNum) * nextCount  # Adjust `nextNum`s to `num`.
      dq.append((num, count))
      while cost > k:  # Remove the rightmost number.
        rightmostNum, rightmostCount = dq.popleft()
        cost -= (rightmostNum - nums[j])
        j -= 1
        if rightmostCount > 1:
          dq.appendleft((rightmostNum, rightmostCount - 1))
      ans += j - i + 1

    return ans

Approach 2: Store indices

• Time: $O(n)$
• Space: $O(n)$

C++JavaPython

class Solution {
 public:
  long long countNonDecreasingSubarrays(vector<int>& nums, int k) {
    long ans = 0;
    long cost = 0;
    // Store indices of `nums` with `nums[i]` in non-increasing order.
    deque<int> dq;

    for (int i = nums.size() - 1, j = nums.size() - 1; i >= 0; --i) {
      const int num = nums[i];
      while (!dq.empty() && nums[dq.back()] < num) {
        const int l = dq.back();
        dq.pop_back();
        const int r = dq.empty() ? j + 1 : dq.back();
        // Adjust `nums[l]` to `num`.
        cost += static_cast<long>(r - l) * (num - nums[l]);
      }
      dq.push_back(i);
      while (cost > k) {  // Remove the rightmost number.
        cost -= nums[dq.front()] - nums[j];
        if (dq.front() == j)
          dq.pop_front();
        --j;
      }
      ans += j - i + 1;
    }

    return ans;
  }
};

class Solution {
  public long countNonDecreasingSubarrays(int[] nums, int k) {
    long ans = 0;
    long cost = 0;
    // Store indices of `nums` with `nums[i]` in non-increasing order.
    Deque<Integer> dq = new ArrayDeque<>();

    for (int i = nums.length - 1, j = nums.length - 1; i >= 0; --i) {
      final int num = nums[i];
      while (!dq.isEmpty() && nums[dq.peekLast()] < num) {
        final int l = dq.pollLast();
        final int r = dq.isEmpty() ? j + 1 : dq.peekLast();
        // Adjust `nums[l]` to `num`.
        cost += (long) (r - l) * (num - nums[l]);
      }
      dq.offerLast(i);
      while (cost > k) { // Remove the rightmost number.
        cost -= nums[dq.peekFirst()] - nums[j];
        if (dq.peekFirst() == j)
          dq.pollFirst();
        --j;
      }
      ans += j - i + 1;
    }

    return ans;
  }
}

class Solution:
  def countNonDecreasingSubarrays(self, nums: list[int], k: int) -> int:
    ans = 0
    cost = 0
    # Store indices (i) of nums with nums[i] in non-increasing order.
    dq = collections.deque()

    j = len(nums) - 1
    for i, num in reversed(list(enumerate(nums))):
      while dq and nums[dq[-1]] < num:
        l = dq.pop()
        r = dq[-1] if dq else j + 1
        cost += (r - l) * (num - nums[l])  # Adjust `nums[l]` to `num`.
      dq.append(i)
      while cost > k:  # Remove the rightmost number.
        cost -= nums[dq[0]] - nums[j]
        if dq[0] == j:
          dq.popleft()
        j -= 1
      ans += j - i + 1

    return ans