Count Non-Decreasing Subarrays After K Operations

1. Initialize a counter count for valid subarrays to 0.
2. Iterate through each possible starting index i from 0 to n-1.
3. For each i, start a nested loop for the ending index j from i to n-1.
4. Inside the inner loop, maintain the currentCost to make the subarray nums[i..j] non-decreasing, and the value of the last element in the modified subarray, prevVal.
5. For the first element of the subarray (j == i), the cost is 0 and prevVal is nums[i].
6. When extending the subarray from j-1 to j, update the cost. If nums[j] is less than prevVal, we must increment nums[j] to prevVal. The cost increases by prevVal - nums[j]. The new prevVal remains the same because the modified element at index j is now prevVal.
7. If nums[j] is not less than prevVal, no cost is incurred at this step, and prevVal is updated to nums[j].
8. After updating the cost for the subarray nums[i..j], check if currentCost <= k.
9. If it is, increment count.
10. If the cost exceeds k, we can break the inner loop for the current i. This is because for any j' > j, the cost for nums[i..j'] will also be greater than k, as costs are non-decreasing when extending the subarray.

We can iterate through all possible subarrays using two nested loops, with i as the starting index and j as the ending index. For each subarray nums[i..j], we can calculate the cost to make it non-decreasing on the fly.

For a fixed starting index i, as we extend the subarray by incrementing j, we can maintain the running cost. Let's say we have the cost for nums[i..j-1] and the last element of its modified version is b_{j-1}. To get the cost for nums[i..j], we add max(0, b_{j-1} - nums[j]). The new last element b_j will be max(b_{j-1}, nums[j]).

An important optimization is that for a fixed i, the cost is a non-decreasing function of j. This means if cost(i, j) > k, then for all j' > j, cost(i, j') will also be greater than k. Thus, we can break the inner loop as soon as the cost exceeds k, avoiding unnecessary computations. Despite this optimization, the overall complexity remains quadratic.

1class Solution {2    public long countSubarrays(int[] nums, int k) {3        int n = nums.length;4        long count = 0;5        for (int i = 0; i < n; i++) {6            long currentCost = 0;7            long prevVal = 0;8            for (int j = i; j < n; j++) {9                if (j == i) {10                    prevVal = nums[j];11                } else {12                    if (nums[j] < prevVal) {13                        currentCost += prevVal - nums[j];14                    } else {15                        prevVal = nums[j];16                    }17                }18                19                if (currentCost <= k) {20                    count++;21                } else {22                    break; // Cost will only increase for this i23                }24            }25        }26        return count;27    }28}

1. Precomputation:
   • Compute next_greater[i]: the index of the first element to the right of i that is greater than nums[i]. This can be done in O(n) using a monotonic stack.
   • Compute prefix_sum array for nums for O(1) range sum queries. Use long to prevent overflow.
   • Build binary lifting tables: up[k][i] stores the 2^k-th next greater element from i, and path_cost[k][i] stores the cost contribution of the b sequence from i to up[k][i]-1. This takes O(n log n).
2. Main Loop:
   • Iterate i from 0 to n-1.
   • For each i, binary search for the maximum j in [i, n-1] such that cost(i, j) <= k.
3. Check Function cost(i, j):
   • This function calculates the cost for nums[i..j] in O(log n) using the precomputed tables.
   • It traverses the next_greater chain from i using the up table (binary lifting) to find all q_k's in [i, j] and sums up their path_cost contributions.
   • It adds the cost for the final segment from the last q_m to j.
   • Finally, it subtracts the sum of original numbers in nums[i..j] (using prefix_sum) to get the total cost.
4. Count Subarrays:
   • If the binary search finds a maximum valid j_{max}, then all subarrays nums[i..j] for j from i to j_{max} are valid. Add j_{max} - i + 1 to the total count.
5. Return the total count.

To make a subarray nums[i..j] non-decreasing with minimum cost, we create a new array b where b_i = nums[i] and b_l = max(nums[l], b_{l-1}) for l > i. The total cost is sum(b_l - nums[l]) for l from i to j.

The sequence b is piecewise constant. The value of b_l only changes from b_{l-1} if nums[l] > b_{l-1}. This means the values in the b sequence are determined by a chain of 'next greater elements' starting from i. Let this chain of indices be q_1=i, q_2, q_3, ... where q_{k+1} is the index of the first element greater than nums[q_k] to its right. The cost can be calculated by summing costs over segments defined by these q indices.

To speed up the cost calculation for an arbitrary (i, j), we can precompute structures using binary lifting. We build a table up[k][i] that gives the 2^k-th next greater element of i. Alongside, a path_cost[k][i] table stores the sum of b_l contributions for the path of 2^k steps. With these tables, we can find the q sequence within [i, j] and calculate the cost in O(log n) time.

For each starting index i, we can then binary search for the maximum j such that cost(i, j) <= k. Since cost(i, j) is monotonic with j, this is feasible. The binary search takes O(log n) iterations, and each check takes O(log n), leading to an O(log^2 n) time for each i.

1// Conceptual structure, not full implementation2class Solution {3    public long countSubarrays(int[] nums, int k) {4        int n = nums.length;5        int logn = (int) (Math.log(n) / Math.log(2)) + 1;6 7        // Precomputation steps (O(n log n))8        int[] nextGreater = new int[n];9        // ... fill nextGreater using a stack (O(n))10 11        long[] prefixSum = new long[n + 1];12        // ... fill prefixSum (O(n))13 14        int[][] up = new int[logn][n];15        long[][] pathCost = new long[logn][n];16        // ... fill binary lifting tables (O(n log n))17 18        long totalCount = 0;19        for (int i = 0; i < n; i++) {20            int low = i, high = n - 1;21            int maxJ = i - 1;22            while (low <= high) {23                int mid = low + (high - low) / 2;24                if (calculateCost(i, mid, nums, prefixSum, up, pathCost, logn) <= k) {25                    maxJ = mid;26                    low = mid + 1;27                } else {28                    high = mid - 1;29                }30            }31            if (maxJ >= i) {32                totalCount += (maxJ - i + 1);33            }34        }35        return totalCount;36    }37 38    private long calculateCost(int i, int j, /*...precomputed tables...*/) {39        // O(log n) cost calculation using binary lifting40        long currentValSum = 0;41        int curr = i;42        // Traverse next_greater chain using 'up' table43        // ... logic ...44        long cost = currentValSum - (prefixSum[j + 1] - prefixSum[i]);45        return cost;46    }47}

1. Precomputation: Same as the O(n log^2 n) approach. Build next_greater, prefix_sum, up, and path_cost tables. This takes O(n log n).
2. Main Loop: Iterate i from 0 to n-1.
3. Greedy Path Traversal: For each i, instead of binary searching for j, we find the maximum reachable j more directly. a. Find Farthest q Node: Greedily traverse the next_greater chain from i using the up table. At each step, try to jump the largest possible distance (2^k) without exceeding the budget k. This finds the farthest node q_m in the chain that can be part of a valid subarray starting at i. This takes O(log n). b. Find Final j: After reaching q_m, we have some remaining budget. The b-value for all subsequent elements will be nums[q_m]. We need to find how many additional elements s we can include, i.e., find max j = q_m + s. The cost for these s+1 elements is (s+1)*nums[q_m] - (prefix_sum[q_m+s+1] - prefix_sum[q_m]). Since this cost is monotonic in s, we can binary search for the maximum s that fits the remaining budget. This also takes O(log n).
4. Count Subarrays: The two steps above give the maximum j_{max} for the starting index i. Add j_{max} - i + 1 to the total count.
5. Return the total count.

The core idea is to improve how we find the maximum valid j for each starting i. Instead of a generic binary search on the index j, we can leverage the structure of the cost function and the next_greater chain.

For a fixed i, we can think of constructing the valid subarray by moving along the next_greater path. We can use our binary lifting up table to make large jumps. We start at curr = i with a budget of k. We greedily try to jump as far as possible. For k from log n down to 0, we check if we can afford the cost of jumping 2^k steps to up[k][curr]. The cost of this segment is pre-calculated in path_cost and prefix_sum. If we can afford it, we make the jump, update curr, and deduct the cost from our budget. This O(log n) process finds the last q node, let's call it last_q, that we can reach.

After reaching last_q, the b-value becomes constant at nums[last_q]. We have some remaining budget. We need to find how many more elements we can append to the subarray. For any j > last_q, the additional cost incurred is (j - last_q + 1) * nums[last_q] - (prefix_sum[j+1] - prefix_sum[last_q]). This cost function is monotonic with j. Therefore, we can perform a second binary search on j in the range [last_q, n-1] to find the maximum j that satisfies the remaining budget. This binary search also takes O(log n).

Since both the greedy traversal and the final binary search take O(log n), the total time to find j_{max} for a given i is O(log n). This brings the overall time complexity down to O(n log n).

1// Conceptual structure for the O(n log n) optimization2class Solution {3    public long countSubarrays(int[] nums, int k) {4        // ... same precomputation as O(n log^2 n) approach ...5 6        long totalCount = 0;7        for (int i = 0; i < n; i++) {8            long currentCost = 0;9            int curr = i;10 11            // Greedily walk the next_greater path (O(log n))12            for (int bit = logn - 1; bit >= 0; bit--) {13                int nextNode = up[bit][curr];14                if (nextNode != n) { // Assuming n is sentinel for no next greater15                    long costToJump = pathCost[bit][curr] - (prefixSum[nextNode] - prefixSum[curr]);16                    if (currentCost + costToJump <= k) {17                        currentCost += costToJump;18                        curr = nextNode;19                    }20                }21            }22 23            // Now curr is the farthest q-node we can reach.24            // Binary search for the final segment length (O(log n))25            int low = curr, high = n - 1, maxJ = curr - 1;26            // ... BS logic to find how many more elements can be added ...27            // The cost check inside BS is: 28            // currentCost + (long)(mid - curr + 1) * nums[curr] - (prefixSum[mid + 1] - prefixSum[curr]) <= k29            30            if (maxJ >= i) {31                totalCount += (maxJ - i + 1);32            }33        }34        return totalCount;35    }36}