Count Non-Decreasing Subarrays After K Operations
Problem
You are given an array nums of n integers and an integer k.
For each subarray of nums, you can apply up to k operations on it. In each operation, you increment any element of the subarray by 1.
Note that each subarray is considered independently, meaning changes made to one subarray do not persist to another.
Return the number of subarrays that you can make non-decreasing after performing at most k operations.
An array is said to be non-decreasing if each element is greater than or equal to its previous element, if it exists.
Examples
Example 1
Input: nums = [6,3,1,2,4,4], k = 7
Output: 17
Explanation:
Out of all 21 possible subarrays of `nums`, only the subarrays `[6, 3, 1]`,
`[6, 3, 1, 2]`, `[6, 3, 1, 2, 4]` and `[6, 3, 1, 2, 4, 4]` cannot be made non-
decreasing after applying up to k = 7 operations. Thus, the number of non-
decreasing subarrays is `21 - 4 = 17`.
Example 2
Input: nums = [6,3,1,3,6], k = 4
Output: 12
Explanation:
The subarray `[3, 1, 3, 6]` along with all subarrays of `nums` with three or
fewer elements, except `[6, 3, 1]`, can be made non-decreasing after `k`
operations. There are 5 subarrays of a single element, 4 subarrays of two
elements, and 2 subarrays of three elements except `[6, 3, 1]`, so there are
`1 + 5 + 4 + 2 = 12` subarrays that can be made non-decreasing.
Constraints
1 <= nums.length <= 10^51 <= nums[i] <= 10^91 <= k <= 10^9
Solution
Method 1 – Sliding Window with Greedy Increment Tracking
Intuition
For each subarray, we want to know if we can make it non-decreasing by incrementing elements at most k times. We can use a sliding window and, for each window, calculate the minimum number of increments needed to make the subarray non-decreasing. If the total increments needed is ≤ k, the subarray is valid.
Approach
- For each possible starting index
i, expand the window to the right as far as possible such that the total increments needed does not exceed k. - For each window, keep track of the required increments to make the subarray non-decreasing:
- If the next element is less than the previous, we need to increment it to match the previous value.
- Accumulate the total increments needed.
- For each valid window, count the subarray.
- Return the total count.
Code
C++
class Solution {
public:
long long countNonDecreasingSubarrays(vector<int>& nums, int k) {
int n = nums.size();
long long ans = 0;
for (int i = 0; i < n; ++i) {
int ops = 0;
int prev = nums[i];
for (int j = i; j < n; ++j) {
if (nums[j] < prev) {
ops += prev - nums[j];
}
if (ops > k) break;
ans++;
prev = max(prev, nums[j]);
}
}
return ans;
}
};
Go
func CountNonDecreasingSubarrays(nums []int, k int) int64 {
n := len(nums)
var ans int64
for i := 0; i < n; i++ {
ops := 0
prev := nums[i]
for j := i; j < n; j++ {
if nums[j] < prev {
ops += prev - nums[j]
}
if ops > k {
break
}
ans++
if nums[j] > prev {
prev = nums[j]
}
}
}
return ans
}
Java
class Solution {
public long countNonDecreasingSubarrays(int[] nums, int k) {
int n = nums.length;
long ans = 0;
for (int i = 0; i < n; i++) {
int ops = 0, prev = nums[i];
for (int j = i; j < n; j++) {
if (nums[j] < prev) {
ops += prev - nums[j];
}
if (ops > k) break;
ans++;
prev = Math.max(prev, nums[j]);
}
}
return ans;
}
}
Kotlin
class Solution {
fun countNonDecreasingSubarrays(nums: IntArray, k: Int): Long {
val n = nums.size
var ans = 0L
for (i in 0 until n) {
var ops = 0
var prev = nums[i]
for (j in i until n) {
if (nums[j] < prev) {
ops += prev - nums[j]
}
if (ops > k) break
ans++
prev = maxOf(prev, nums[j])
}
}
return ans
}
}
Python
class Solution:
def countNonDecreasingSubarrays(self, nums: list[int], k: int) -> int:
n = len(nums)
ans = 0
for i in range(n):
ops = 0
prev = nums[i]
for j in range(i, n):
if nums[j] < prev:
ops += prev - nums[j]
if ops > k:
break
ans += 1
prev = max(prev, nums[j])
return ans
Rust
impl Solution {
pub fn count_non_decreasing_subarrays(nums: Vec<i32>, k: i32) -> i64 {
let n = nums.len();
let mut ans = 0i64;
for i in 0..n {
let mut ops = 0;
let mut prev = nums[i];
for j in i..n {
if nums[j] < prev {
ops += prev - nums[j];
}
if ops > k {
break;
}
ans += 1;
prev = prev.max(nums[j]);
}
}
ans
}
}
TypeScript
class Solution {
countNonDecreasingSubarrays(nums: number[], k: number): number {
const n = nums.length;
let ans = 0;
for (let i = 0; i < n; i++) {
let ops = 0, prev = nums[i];
for (let j = i; j < n; j++) {
if (nums[j] < prev) {
ops += prev - nums[j];
}
if (ops > k) break;
ans++;
prev = Math.max(prev, nums[j]);
}
}
return ans;
}
}
Complexity
- ⏰ Time complexity:
O(n^2), since we check all subarrays. - 🧺 Space complexity:
O(1), only constant extra space is used.
Method 2 – Sliding Window with Next Greater Index (NGI)
Intuition
If a subarray [L, R] can be made non-decreasing with at most k operations, then any subarray [L, X] where L ≤ X ≤ R can also be made non-decreasing. For a subarray, the minimum operations needed to make all elements equal to the largest value T is (R - L + 1) * T - sum(L, R). This linear relationship allows efficient window adjustment.
Approach
- Precompute the next greater index (
ngi) for each element using a stack. - Use a sliding window
[idx, i]to track the current valid subarray. - Maintain the largest value in the window (
curLargeValue) and the total operations needed (totalK). - If
totalKexceedsk, adjust the window usingngito efficiently skip invalid regions. - For each valid window, count the subarrays ending at
i. - After processing, add the remaining valid subarrays.
Code
C++
#define ll long long
class Solution {
public:
long long countNonDecreasingSubarrays(vector<int>& nums, int k) {
int n = nums.size();
stack<int> st;
vector<int> ngi(n, n);
for (int i = 0; i < n; i++) {
while (!st.empty() && nums[i] > nums[st.top()]) {
ngi[st.top()] = i;
st.pop();
}
st.push(i);
}
int idx = 0;
ll count = 0;
int curLargeValue = nums[idx];
ll totalK = 0;
for (int i = 0; i < n; i++) {
if (nums[i] > curLargeValue) {
curLargeValue = nums[i];
} else {
totalK += curLargeValue - nums[i];
}
while (totalK > k) {
count += i - idx;
if (nums[idx + 1] >= nums[idx]) {
idx++;
} else {
ll prev = nums[idx];
ll prevIdx = idx;
ll cur = nums[idx + 1];
ll curIdx = idx + 1;
curLargeValue = cur;
while (ngi[curIdx] <= i && ngi[prevIdx] != ngi[curIdx]) {
totalK -= (prev - cur) * (ngi[curIdx] - curIdx);
curIdx = ngi[curIdx];
cur = nums[curIdx];
curLargeValue = cur;
}
if (i >= ngi[curIdx] && ngi[prevIdx] == ngi[curIdx]) {
totalK -= (prev - cur) * (ngi[curIdx] - curIdx);
curLargeValue = nums[ngi[curIdx]];
cur = curLargeValue;
curIdx = ngi[curIdx];
prev = cur;
prevIdx = curIdx;
while (i >= ngi[curIdx] && ngi[prevIdx] == ngi[curIdx]) {
curLargeValue = nums[ngi[curIdx]];
curIdx = ngi[curIdx];
prev = cur;
prevIdx = curIdx;
}
} else {
totalK -= (prev - cur) * (i - curIdx + 1);
}
idx++;
}
}
}
count += 1ll * (n - idx) * (n - idx + 1) / 2;
return count;
}
};
Go
func countNonDecreasingSubarrays(nums []int, k int) int64 {
n := len(nums)
ngi := make([]int, n)
for i := range ngi {
ngi[i] = n
}
st := []int{}
for i := 0; i < n; i++ {
for len(st) > 0 && nums[i] > nums[st[len(st)-1]] {
ngi[st[len(st)-1]] = i
st = st[:len(st)-1]
}
st = append(st, i)
}
idx := 0
count := int64(0)
curLargeValue := nums[idx]
totalK := int64(0)
for i := 0; i < n; i++ {
if nums[i] > curLargeValue {
curLargeValue = nums[i]
} else {
totalK += int64(curLargeValue - nums[i])
}
for totalK > int64(k) {
count += int64(i - idx)
if nums[idx+1] >= nums[idx] {
idx++
} else {
prev := nums[idx]
prevIdx := idx
cur := nums[idx+1]
curIdx := idx+1
curLargeValue = cur
for ngi[curIdx] <= i && ngi[prevIdx] != ngi[curIdx] {
totalK -= int64(prev-cur) * int64(ngi[curIdx]-curIdx)
curIdx = ngi[curIdx]
cur = nums[curIdx]
curLargeValue = cur
}
if i >= ngi[curIdx] && ngi[prevIdx] == ngi[curIdx] {
totalK -= int64(prev-cur) * int64(ngi[curIdx]-curIdx)
curLargeValue = nums[ngi[curIdx]]
cur = curLargeValue
curIdx = ngi[curIdx]
prev = cur
prevIdx = curIdx
for i >= ngi[curIdx] && ngi[prevIdx] == ngi[curIdx] {
curLargeValue = nums[ngi[curIdx]]
curIdx = ngi[curIdx]
prev = cur
prevIdx = curIdx
}
} else {
totalK -= int64(prev-cur) * int64(i-curIdx+1)
}
idx++
}
}
}
count += int64(n-idx) * int64(n-idx+1) / 2
return count
}
Java
class Solution {
public long countNonDecreasingSubarrays(int[] nums, int k) {
int n = nums.length;
int[] ngi = new int[n];
Arrays.fill(ngi, n);
Deque<Integer> st = new ArrayDeque<>();
for (int i = 0; i < n; i++) {
while (!st.isEmpty() && nums[i] > nums[st.peek()]) {
ngi[st.pop()] = i;
}
st.push(i);
}
int idx = 0;
long count = 0;
int curLargeValue = nums[idx];
long totalK = 0;
for (int i = 0; i < n; i++) {
if (nums[i] > curLargeValue) {
curLargeValue = nums[i];
} else {
totalK += curLargeValue - nums[i];
}
while (totalK > k) {
count += i - idx;
if (nums[idx + 1] >= nums[idx]) {
idx++;
} else {
long prev = nums[idx];
int prevIdx = idx;
long cur = nums[idx + 1];
int curIdx = idx + 1;
curLargeValue = (int)cur;
while (ngi[curIdx] <= i && ngi[prevIdx] != ngi[curIdx]) {
totalK -= (prev - cur) * (ngi[curIdx] - curIdx);
curIdx = ngi[curIdx];
cur = nums[curIdx];
curLargeValue = (int)cur;
}
if (i >= ngi[curIdx] && ngi[prevIdx] == ngi[curIdx]) {
totalK -= (prev - cur) * (ngi[curIdx] - curIdx);
curLargeValue = nums[ngi[curIdx]];
cur = curLargeValue;
curIdx = ngi[curIdx];
prev = cur;
prevIdx = curIdx;
while (i >= ngi[curIdx] && ngi[prevIdx] == ngi[curIdx]) {
curLargeValue = nums[ngi[curIdx]];
curIdx = ngi[curIdx];
prev = cur;
prevIdx = curIdx;
}
} else {
totalK -= (prev - cur) * (i - curIdx + 1);
}
idx++;
}
}
}
count += 1L * (n - idx) * (n - idx + 1) / 2;
return count;
}
}
Kotlin
class Solution {
fun countNonDecreasingSubarrays(nums: IntArray, k: Int): Long {
val n = nums.size
val ngi = IntArray(n) { n }
val st = ArrayDeque<Int>()
for (i in 0 until n) {
while (st.isNotEmpty() && nums[i] > nums[st.last()]) {
ngi[st.removeLast()] = i
}
st.addLast(i)
}
var idx = 0
var count = 0L
var curLargeValue = nums[idx]
var totalK = 0L
for (i in 0 until n) {
if (nums[i] > curLargeValue) {
curLargeValue = nums[i]
} else {
totalK += curLargeValue - nums[i]
}
while (totalK > k) {
count += i - idx
if (nums[idx + 1] >= nums[idx]) {
idx++
} else {
var prev = nums[idx].toLong()
var prevIdx = idx
var cur = nums[idx + 1].toLong()
var curIdx = idx + 1
curLargeValue = cur.toInt()
while (ngi[curIdx] <= i && ngi[prevIdx] != ngi[curIdx]) {
totalK -= (prev - cur) * (ngi[curIdx] - curIdx)
curIdx = ngi[curIdx]
cur = nums[curIdx].toLong()
curLargeValue = cur.toInt()
}
if (i >= ngi[curIdx] && ngi[prevIdx] == ngi[curIdx]) {
totalK -= (prev - cur) * (ngi[curIdx] - curIdx)
curLargeValue = nums[ngi[curIdx]]
cur = curLargeValue.toLong()
curIdx = ngi[curIdx]
prev = cur
prevIdx = curIdx
while (i >= ngi[curIdx] && ngi[prevIdx] == ngi[curIdx]) {
curLargeValue = nums[ngi[curIdx]]
curIdx = ngi[curIdx]
prev = cur
prevIdx = curIdx
}
} else {
totalK -= (prev - cur) * (i - curIdx + 1)
}
idx++
}
}
}
count += 1L * (n - idx) * (n - idx + 1) / 2
return count
}
}
Python
class Solution:
def countNonDecreasingSubarrays(self, nums: list[int], k: int) -> int:
n = len(nums)
ngi = [n] * n
st = []
for i in range(n):
while st and nums[i] > nums[st[-1]]:
ngi[st.pop()] = i
st.append(i)
idx = 0
count = 0
curLargeValue = nums[idx]
totalK = 0
for i in range(n):
if nums[i] > curLargeValue:
curLargeValue = nums[i]
else:
totalK += curLargeValue - nums[i]
while totalK > k:
count += i - idx
if nums[idx + 1] >= nums[idx]:
idx += 1
else:
prev = nums[idx]
prevIdx = idx
cur = nums[idx + 1]
curIdx = idx + 1
curLargeValue = cur
while ngi[curIdx] <= i && ngi[prevIdx] != ngi[curIdx]:
totalK -= (prev - cur) * (ngi[curIdx] - curIdx)
curIdx = ngi[curIdx]
cur = nums[curIdx]
curLargeValue = cur
if i >= ngi[curIdx] && ngi[prevIdx] == ngi[curIdx]:
totalK -= (prev - cur) * (ngi[curIdx] - curIdx)
curLargeValue = nums[ngi[curIdx]]
cur = curLargeValue
curIdx = ngi[curIdx]
prev = cur
prevIdx = curIdx
while i >= ngi[curIdx] && ngi[prevIdx] == ngi[curIdx]:
curLargeValue = nums[ngi[curIdx]]
curIdx = ngi[curIdx]
prev = cur
prevIdx = curIdx
else:
totalK -= (prev - cur) * (i - curIdx + 1)
idx += 1
count += (n - idx) * (n - idx + 1) // 2
return count
Rust
impl Solution {
pub fn count_non_decreasing_subarrays(nums: Vec<i32>, k: i32) -> i64 {
let n = nums.len();
let mut ngi = vec![n; n];
let mut st = Vec::new();
for i in 0..n {
while let Some(&top) = st.last() {
if nums[i] > nums[top] {
ngi[top] = i;
st.pop();
} else {
break;
}
}
st.push(i);
}
let mut idx = 0;
let mut count = 0i64;
let mut cur_large = nums[idx];
let mut total_k = 0i64;
for i in 0..n {
if nums[i] > cur_large {
cur_large = nums[i];
} else {
total_k += (cur_large - nums[i]) as i64;
}
while total_k > k as i64 {
count += (i - idx) as i64;
if nums[idx + 1] >= nums[idx] {
idx += 1;
} else {
let mut prev = nums[idx] as i64;
let mut prev_idx = idx;
let mut cur = nums[idx + 1] as i64;
let mut cur_idx = idx + 1;
cur_large = cur as i32;
while ngi[cur_idx] <= i && ngi[prev_idx] != ngi[cur_idx] {
total_k -= (prev - cur) * (ngi[cur_idx] as i64 - cur_idx as i64);
cur_idx = ngi[cur_idx];
cur = nums[cur_idx] as i64;
cur_large = cur as i32;
}
if i >= ngi[cur_idx] && ngi[prev_idx] == ngi[cur_idx] {
total_k -= (prev - cur) * (ngi[cur_idx] as i64 - cur_idx as i64);
cur_large = nums[ngi[cur_idx]];
cur = cur_large as i64;
cur_idx = ngi[cur_idx];
prev = cur;
prev_idx = cur_idx;
while i >= ngi[cur_idx] && ngi[prev_idx] == ngi[cur_idx] {
cur_large = nums[ngi[cur_idx]];
cur_idx = ngi[cur_idx];
prev = cur;
prev_idx = cur_idx;
}
} else {
total_k -= (prev - cur) * ((i - cur_idx + 1) as i64);
}
idx += 1;
}
}
}
count += ((n - idx) as i64) * ((n - idx + 1) as i64) / 2;
count
}
}
TypeScript
class Solution {
countNonDecreasingSubarrays(nums: number[], k: number): number {
const n = nums.length;
const ngi: number[] = Array(n).fill(n);
const st: number[] = [];
for (let i = 0; i < n; i++) {
while (st.length && nums[i] > nums[st[st.length - 1]]) {
ngi[st.pop()!] = i;
}
st.push(i);
}
let idx = 0;
let count = 0;
let curLargeValue = nums[idx];
let totalK = 0;
for (let i = 0; i < n; i++) {
if (nums[i] > curLargeValue) {
curLargeValue = nums[i];
} else {
totalK += curLargeValue - nums[i];
}
while (totalK > k) {
count += i - idx;
if (nums[idx + 1] >= nums[idx]) {
idx++;
} else {
let prev = nums[idx];
let prevIdx = idx;
let cur = nums[idx + 1];
let curIdx = idx + 1;
curLargeValue = cur;
while (ngi[curIdx] <= i && ngi[prevIdx] != ngi[curIdx]) {
totalK -= (prev - cur) * (ngi[curIdx] - curIdx);
curIdx = ngi[curIdx];
cur = nums[curIdx];
curLargeValue = cur;
}
if (i >= ngi[curIdx] && ngi[prevIdx] == ngi[curIdx]) {
totalK -= (prev - cur) * (ngi[curIdx] - curIdx);
curLargeValue = nums[ngi[curIdx]];
cur = curLargeValue;
curIdx = ngi[curIdx];
prev = cur;
prevIdx = curIdx;
while (i >= ngi[curIdx] && ngi[prevIdx] == ngi[curIdx]) {
curLargeValue = nums[ngi[curIdx]];
curIdx = ngi[curIdx];
prev = cur;
prevIdx = curIdx;
}
} else {
totalK -= (prev - cur) * (i - curIdx + 1);
}
idx++;
}
}
}
count += (n - idx) * (n - idx + 1) / 2;
return count;
}
}
Complexity
- ⏰ Time complexity:
O(n), since each element is pushed and popped from the stack at most once and the sliding window is amortized. - 🧺 Space complexity:
O(n), for the stack andngiarray.