402. Minimum Operations to Make Elements Within K Subarrays Equal

Equalizing a window costs least when every element moves to the window's median; precompute that cost for each window, then DP to pick k non-overlapping windows of minimum total cost.

Two sub-problems

The task splits cleanly into two parts: (1) for every possible length-x window, how cheaply can we make it internally equal, and (2) which k non-overlapping windows should we actually use.

Part 1 — cost of one window

Inside a single window we may pick any common target value t; the cost is the sum of |a - t| over the window. That sum is minimized when t is the median of the window (moving the target off the median helps one side but hurts the other side by at least as much). So:

cost(window) = sum of |a_i - median|

Sort the window, let m = x // 2 be the median index, and split the sum into the part below the median (lift it up) and the part above (pull it down). Using prefix sums of the sorted window this is O(x) per window, and as the window slides by one we only remove one element and insert one, keeping a sorted copy with binary search.

Part 2 — choose k windows with a DP

Now we have cost[s] for each start s. We want k non-overlapping windows (a window at s occupies indices s .. s+x-1) with the smallest total cost. Sweep left to right with:

dp_j[i] = minimum operations to place j windows entirely inside the first i elements.

At index i we either skip it (dp_j[i] = dp_j[i-1]) or close a window that occupies [i-x, i-1], which starts at i-x and leaves the first i-x elements for the earlier windows:

dp_j[i] = min(dp_j[i-1], dp_{j-1}[i-x] + cost[i-x])

The answer is dp_k[n].

Worked example

nums = [5,-2,1,3,7,3,6,4,-1], x = 3, k = 2.

1. Slide a window of size 3 and compute each window's median cost. The window [-2,1,3] (indices 1..3) has median 1, costing |-2-1| + |1-1| + |3-1| = 5. The window [3,6,4] (indices 5..7) has median 4, costing |3-4| + |6-4| + |4-4| = 3.
2. The DP looks for two non-overlapping windows with the smallest total. Picking starts 1 and 5 (they do not overlap) gives 5 + 3 = 8, which is the minimum.
3. Answer: 8 — matching turning indices 1..3 into [1,1,1] and indices 5..7 into [4,4,4].

Complexity

• Time: O(n · x) to precompute every window cost with the sliding sorted window, plus O(n · k) for the DP. (A balanced two-heap / ordered multiset lowers the cost pass to O(n log x).)
• Space: O(n) for the cost array and the two DP rows.

import java.util.ArrayList;
import java.util.Collections;

class Solution {
    public long minOperations(int[] nums, int x, int k) {
        int n = nums.length;

        // Step 1: cost[s] = minimum ops to make the length-x window starting at
        // s all equal. The cheapest common target is the window's median, so the
        // cost is the sum of absolute deviations from that median. We keep a
        // sorted copy of the current window and read the cost off its prefix sum.
        ArrayList<Integer> win = new ArrayList<>();
        for (int i = 0; i < x; i++) win.add(nums[i]);
        Collections.sort(win);
        long[] prefix = new long[x + 1];
        long[] cost = new long[n - x + 1];
        cost[0] = windowCost(win, prefix, x);

        for (int i = x; i < n; i++) {
            win.remove(Integer.valueOf(nums[i - x]));         // drop element leaving
            int pos = Collections.binarySearch(win, nums[i]); // add element entering
            if (pos < 0) pos = -pos - 1;
            win.add(pos, nums[i]);
            cost[i - x + 1] = windowCost(win, prefix, x);
        }

        // Step 2: choose k non-overlapping windows with minimum total cost.
        // dp after j windows: min ops using j windows inside the first i elements.
        // A window occupying [i-x, i-1] starts at i-x.
        long INF = Long.MAX_VALUE / 4;
        long[] dpPrev = new long[n + 1];                      // j = 0 windows -> 0 ops
        for (int j = 0; j < k; j++) {
            long[] dpCur = new long[n + 1];
            dpCur[0] = INF;
            for (int i = 1; i <= n; i++) {
                long best = dpCur[i - 1];                     // skip index i-1
                if (i >= x && dpPrev[i - x] < INF) {
                    best = Math.min(best, dpPrev[i - x] + cost[i - x]);
                }
                dpCur[i] = best;
            }
            dpPrev = dpCur;
        }
        return dpPrev[n];
    }

    private long windowCost(ArrayList<Integer> win, long[] prefix, int x) {
        for (int i = 0; i < x; i++) prefix[i + 1] = prefix[i] + win.get(i);
        int m = x / 2;                        // median index inside the sorted window
        long median = win.get(m);
        long left = median * m - prefix[m];
        long right = (prefix[x] - prefix[m]) - median * (x - m);
        return left + right;
    }
}