3422. Minimum Operations to Make Subarray Elements Equal 🔒

MediumArrayHash TableMathSliding WindowHeap (Priority Queue)

Problem Description

You are given an integer array nums and an integer k. You can perform operations to increase or decrease any element of nums by 1, and you can do this any number of times.

Your goal is to find the minimum number of operations needed to make at least one subarray of size k have all its elements equal.

For example, if nums = [1, 4, 2, 6] and k = 3, you need to find a consecutive subarray of length 3 (like [1, 4, 2] or [4, 2, 6]) and determine how many operations it would take to make all elements in that subarray equal. You want to choose the subarray that requires the fewest operations.

The key insight is that to minimize operations, you should change all elements in a subarray to their median value. This is because the median minimizes the sum of absolute differences. The solution uses two ordered sets to efficiently track and update the median as it slides through different subarrays of size k. The left set l stores the smaller half of elements, the right set r stores the larger half, and the smallest element in r represents the median. The minimum operations for each subarray equals the sum of differences between each element and the median.

Quick Interview Experience

Help others by sharing your interview experience

Have you seen this problem before?

Intuition

When we need to make all elements in an array equal, the optimal target value is the median. Why? Because the median minimizes the total distance (sum of absolute differences) to all other elements. If we choose any value other than the median, we'd need more operations.

Since we need to check every possible subarray of size k, a naive approach would be to sort each subarray to find its median, then calculate the operations needed. But this would be inefficient, requiring O(k log k) time for each of the n-k+1 subarrays.

The key insight is that as we slide from one subarray to the next, we're essentially removing one element and adding another. Instead of recalculating everything from scratch, we can maintain the median dynamically.

To efficiently maintain a median while adding and removing elements, we can use two balanced sets:

A "left" set l containing the smaller half of elements
A "right" set r containing the larger half of elements

We maintain the invariant that r has either the same number of elements as l or one more. This way, the minimum element in r is always our median.

When sliding the window:

We add the new element to the appropriate set
We rebalance if needed (ensuring r doesn't have more than one extra element compared to l)
We remove the outgoing element from whichever set contains it

The number of operations for any subarray equals: (median * count_of_smaller_elements - sum_of_smaller_elements) + (sum_of_larger_elements - median * count_of_larger_elements). This formula represents moving all smaller elements up to the median and all larger elements down to the median.

Learn more about Math, Sliding Window and Heap (Priority Queue) patterns.

Solution Approach

The implementation uses two SortedList data structures (ordered sets) to maintain the median dynamically as we slide through subarrays of size k.

Data Structures:

l: SortedList storing the smaller half of elements in the current window
r: SortedList storing the larger half of elements (the minimum of r is the median)
s1: Sum of all elements in l
s2: Sum of all elements in r

Algorithm Steps:

Building the Initial Window:
- For each element x in nums, we first add it to the left set l and update s1
- We then pop the maximum element from l and move it to r to maintain balance
- This ensures elements are properly distributed between the two sets
Maintaining Balance:
- After adding to r, we check if len(r) - len(l) > 1
- If true, we move the minimum element from r back to l
- This maintains the invariant that r has at most one more element than l
Calculating Operations for Each Window:
- Once we have k elements (when i >= k - 1), we calculate the operations needed
- The median is r[0] (minimum element in the right set)
- Operations = s2 - r[0] * len(r) + r[0] * len(l) - s1
  - r[0] * len(l) - s1: Operations to move all elements in l up to the median
  - s2 - r[0] * len(r): Operations to move all elements in r down to the median
Sliding the Window:
- Remove the leftmost element of the previous window (nums[j] where j = i - k + 1)
- Check which set contains this element and remove it, updating the corresponding sum
- The process continues with the next iteration automatically adding the new element
Tracking the Minimum:
- We maintain ans to track the minimum operations across all windows
- Return ans after processing all possible subarrays of size k

The time complexity is O(n log k) where each insertion, deletion, and finding min/max in SortedList takes O(log k) time, and we perform these operations for each of the n elements.

Ready to land your dream job?

Unlock your dream job with a 5-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

Let's trace through the algorithm with nums = [1, 4, 2, 6] and k = 3.

Initial Setup:

l = [] (left set for smaller elements)
r = [] (right set for larger elements)
s1 = 0 (sum of left set)
s2 = 0 (sum of right set)
ans = infinity

Step 1: Process nums[0] = 1

Add 1 to l: l = [1], s1 = 1
Move max of l (which is 1) to r: l = [], s1 = 0, r = [1], s2 = 1
Check balance: len(r) - len(l) = 1 - 0 = 1 (balanced)
Window size = 1 < 3, so continue

Step 2: Process nums[1] = 4

Add 4 to l: l = [4], s1 = 4
Move max of l (which is 4) to r: l = [], s1 = 0, r = [1, 4], s2 = 5
Check balance: len(r) - len(l) = 2 - 0 = 2 > 1 (unbalanced!)
Move min of r (which is 1) to l: l = [1], s1 = 1, r = [4], s2 = 4
Window size = 2 < 3, so continue

Step 3: Process nums[2] = 2

Add 2 to l: l = [1, 2], s1 = 3
Move max of l (which is 2) to r: l = [1], s1 = 1, r = [2, 4], s2 = 6
Check balance: len(r) - len(l) = 2 - 1 = 1 (balanced)
Window size = 3 = k, calculate operations!
- Window is [1, 4, 2]
- Median = r[0] = 2
- Operations = s2 - 2*len(r) + 2*len(l) - s1
- Operations = 6 - 2*2 + 2*1 - 1 = 6 - 4 + 2 - 1 = 3
- Update ans = 3
Remove leftmost element (nums[0] = 1) from window:
- 1 is in l, so remove it: l = [], s1 = 0

Step 4: Process nums[3] = 6

Add 6 to l: l = [6], s1 = 6
Move max of l (which is 6) to r: l = [], s1 = 0, r = [2, 4, 6], s2 = 12
Check balance: len(r) - len(l) = 3 - 0 = 3 > 1 (unbalanced!)
Move min of r (which is 2) to l: l = [2], s1 = 2, r = [4, 6], s2 = 10
Window size = 3 = k, calculate operations!
- Window is [4, 2, 6]
- Median = r[0] = 4
- Operations = s2 - 4*len(r) + 4*len(l) - s1
- Operations = 10 - 4*2 + 4*1 - 2 = 10 - 8 + 4 - 2 = 4
- ans = min(3, 4) = 3

Result: The minimum operations needed is 3. This occurs for the subarray [1, 4, 2], where we need:

1 operation to change 1 to 2
2 operations to change 4 to 2
0 operations for 2 (already at median)
Total: 3 operations

Solution Implementation

1class Solution:
2    def minOperations(self, nums: List[int], k: int) -> int:
3        from sortedcontainers import SortedList
4        import math
5      
6        # Two sorted lists to maintain elements below and above median
7        # left_half contains smaller elements, right_half contains larger elements
8        left_half = SortedList()
9        right_half = SortedList()
10      
11        # Sum of elements in left_half and right_half respectively
12        left_sum = 0
13        right_sum = 0
14      
15        # Initialize answer to infinity
16        min_cost = math.inf
17      
18        for i, current_num in enumerate(nums):
19            # Add current number to left_half first
20            left_half.add(current_num)
21            left_sum += current_num
22          
23            # Move the largest element from left_half to right_half
24            # This helps maintain the median property
25            largest_left = left_half.pop()
26            left_sum -= largest_left
27            right_half.add(largest_left)
28            right_sum += largest_left
29          
30            # Balance the two halves: right_half should have at most one more element than left_half
31            # This ensures the median is always the smallest element in right_half
32            if len(right_half) - len(left_half) > 1:
33                smallest_right = right_half.pop(0)
34                right_sum -= smallest_right
35                left_half.add(smallest_right)
36                left_sum += smallest_right
37          
38            # Once we have at least k elements, calculate the cost for current window
39            if i >= k - 1:
40                # The median is the smallest element in right_half
41                median = right_half[0]
42              
43                # Cost to make all elements equal to median:
44                # - Elements smaller than median need to be increased: median * count - sum
45                # - Elements larger than median need to be decreased: sum - median * count
46                cost = (right_sum - median * len(right_half)) + (median * len(left_half) - left_sum)
47                min_cost = min(min_cost, cost)
48              
49                # Remove the leftmost element of the sliding window
50                window_start = i - k + 1
51                element_to_remove = nums[window_start]
52              
53                # Remove from appropriate half and update sum
54                if element_to_remove in right_half:
55                    right_half.remove(element_to_remove)
56                    right_sum -= element_to_remove
57                else:
58                    left_half.remove(element_to_remove)
59                    left_sum -= element_to_remove
60      
61        return min_cost
62

1class Solution {
2    public long minOperations(int[] nums, int k) {
3        // TreeMap to store the smaller half of elements (max heap behavior)
4        TreeMap<Integer, Integer> smallerHalf = new TreeMap<>();
5        // TreeMap to store the larger half of elements (min heap behavior)
6        TreeMap<Integer, Integer> largerHalf = new TreeMap<>();
7      
8        // Sum of elements in smaller half
9        long smallerSum = 0;
10        // Sum of elements in larger half
11        long largerSum = 0;
12        // Count of elements in smaller half
13        int smallerCount = 0;
14        // Count of elements in larger half
15        int largerCount = 0;
16      
17        long minOperations = Long.MAX_VALUE;
18      
19        for (int i = 0; i < nums.length; i++) {
20            // Add current element to smaller half initially
21            smallerHalf.merge(nums[i], 1, Integer::sum);
22            smallerSum += nums[i];
23            smallerCount++;
24          
25            // Move the largest element from smaller half to larger half
26            int maxFromSmaller = smallerHalf.lastKey();
27            if (smallerHalf.merge(maxFromSmaller, -1, Integer::sum) == 0) {
28                smallerHalf.remove(maxFromSmaller);
29            }
30            smallerSum -= maxFromSmaller;
31            smallerCount--;
32          
33            largerHalf.merge(maxFromSmaller, 1, Integer::sum);
34            largerSum += maxFromSmaller;
35            largerCount++;
36          
37            // Balance the two halves: larger half should have at most one more element than smaller half
38            if (largerCount - smallerCount > 1) {
39                int minFromLarger = largerHalf.firstKey();
40                if (largerHalf.merge(minFromLarger, -1, Integer::sum) == 0) {
41                    largerHalf.remove(minFromLarger);
42                }
43                largerSum -= minFromLarger;
44                largerCount--;
45              
46                smallerHalf.merge(minFromLarger, 1, Integer::sum);
47                smallerSum += minFromLarger;
48                smallerCount++;
49            }
50          
51            // When we have at least k elements, calculate the cost for current window
52            if (i >= k - 1) {
53                // The median is the smallest element in the larger half
54                int median = largerHalf.firstKey();
55                // Cost = sum of (median - smaller elements) + sum of (larger elements - median)
56                // Which simplifies to: (largerSum - median * largerCount) + (median * smallerCount - smallerSum)
57                long currentCost = largerSum - (long)median * largerCount + (long)median * smallerCount - smallerSum;
58                minOperations = Math.min(minOperations, currentCost);
59              
60                // Remove the leftmost element from the window
61                int leftmostIndex = i - k + 1;
62                int elementToRemove = nums[leftmostIndex];
63              
64                // Remove from the appropriate half
65                if (largerHalf.containsKey(elementToRemove)) {
66                    if (largerHalf.merge(elementToRemove, -1, Integer::sum) == 0) {
67                        largerHalf.remove(elementToRemove);
68                    }
69                    largerSum -= elementToRemove;
70                    largerCount--;
71                } else {
72                    if (smallerHalf.merge(elementToRemove, -1, Integer::sum) == 0) {
73                        smallerHalf.remove(elementToRemove);
74                    }
75                    smallerSum -= elementToRemove;
76                    smallerCount--;
77                }
78            }
79        }
80      
81        return minOperations;
82    }
83}
84

1class Solution {
2public:
3    long long minOperations(vector<int>& nums, int k) {
4        // Two multisets to maintain the smaller and larger halves of the sliding window
5        // leftHalf contains smaller elements, rightHalf contains larger elements
6        multiset<int> leftHalf, rightHalf;
7        long long leftSum = 0, rightSum = 0;  // Sum of elements in each half
8        long long minCost = 1e18;  // Initialize minimum cost to a large value
9      
10        for (int i = 0; i < nums.size(); ++i) {
11            // Step 1: Add current element to the left half initially
12            leftHalf.insert(nums[i]);
13            leftSum += nums[i];
14          
15            // Step 2: Move the largest element from left half to right half
16            int largestFromLeft = *leftHalf.rbegin();
17            leftHalf.erase(leftHalf.find(largestFromLeft));
18            leftSum -= largestFromLeft;
19            rightHalf.insert(largestFromLeft);
20            rightSum += largestFromLeft;
21          
22            // Step 3: Balance the two halves - ensure their sizes differ by at most 1
23            // If right half has more than 1 extra element, move smallest to left
24            if (rightHalf.size() - leftHalf.size() > 1) {
25                int smallestFromRight = *rightHalf.begin();
26                rightHalf.erase(rightHalf.find(smallestFromRight));
27                rightSum -= smallestFromRight;
28                leftHalf.insert(smallestFromRight);
29                leftSum += smallestFromRight;
30            }
31          
32            // Step 4: Process window when we have at least k elements
33            if (i >= k - 1) {
34                // The median is the smallest element in the right half
35                long long median = *rightHalf.begin();
36              
37                // Calculate cost to make all elements equal to median
38                // Cost = (elements greater than median) - (elements less than median)
39                long long currentCost = rightSum - median * (int)rightHalf.size() + 
40                                       median * (int)leftHalf.size() - leftSum;
41                minCost = min(minCost, currentCost);
42              
43                // Remove the leftmost element of the sliding window
44                int elementToRemove = nums[i - k + 1];
45                if (rightHalf.contains(elementToRemove)) {
46                    rightHalf.erase(rightHalf.find(elementToRemove));
47                    rightSum -= elementToRemove;
48                } else {
49                    leftHalf.erase(leftHalf.find(elementToRemove));
50                    leftSum -= elementToRemove;
51                }
52            }
53        }
54      
55        return minCost;
56    }
57};
58

1function minOperations(nums: number[], k: number): number {
2    // Two multisets to maintain the smaller and larger halves of the sliding window
3    // leftHalf contains smaller elements, rightHalf contains larger elements
4    const leftHalf = new Multiset<number>();
5    const rightHalf = new Multiset<number>();
6    let leftSum = 0;   // Sum of elements in left half
7    let rightSum = 0;  // Sum of elements in right half
8    let minCost = Number.MAX_SAFE_INTEGER;  // Initialize minimum cost to a large value
9  
10    for (let i = 0; i < nums.length; i++) {
11        // Step 1: Add current element to the left half initially
12        leftHalf.add(nums[i]);
13        leftSum += nums[i];
14      
15        // Step 2: Move the largest element from left half to right half
16        const largestFromLeft = leftHalf.max()!;
17        leftHalf.delete(largestFromLeft);
18        leftSum -= largestFromLeft;
19        rightHalf.add(largestFromLeft);
20        rightSum += largestFromLeft;
21      
22        // Step 3: Balance the two halves - ensure their sizes differ by at most 1
23        // If right half has more than 1 extra element, move smallest to left
24        if (rightHalf.size - leftHalf.size > 1) {
25            const smallestFromRight = rightHalf.min()!;
26            rightHalf.delete(smallestFromRight);
27            rightSum -= smallestFromRight;
28            leftHalf.add(smallestFromRight);
29            leftSum += smallestFromRight;
30        }
31      
32        // Step 4: Process window when we have at least k elements
33        if (i >= k - 1) {
34            // The median is the smallest element in the right half
35            const median = rightHalf.min()!;
36          
37            // Calculate cost to make all elements equal to median
38            // Cost = (elements greater than median) - (elements less than median)
39            const currentCost = rightSum - median * rightHalf.size + 
40                              median * leftHalf.size - leftSum;
41            minCost = Math.min(minCost, currentCost);
42          
43            // Remove the leftmost element of the sliding window
44            const elementToRemove = nums[i - k + 1];
45            if (rightHalf.has(elementToRemove)) {
46                rightHalf.delete(elementToRemove);
47                rightSum -= elementToRemove;
48            } else {
49                leftHalf.delete(elementToRemove);
50                leftSum -= elementToRemove;
51            }
52        }
53    }
54  
55    return minCost;
56}
57
58// Multiset implementation for TypeScript (similar to C++ multiset)
59class Multiset<T> {
60    private map: Map<T, number> = new Map();
61    private _size: number = 0;
62    private sorted: T[] = [];
63    private needsSort: boolean = false;
64  
65    add(value: T): void {
66        this.map.set(value, (this.map.get(value) || 0) + 1);
67        this._size++;
68        this.needsSort = true;
69    }
70  
71    delete(value: T): boolean {
72        const count = this.map.get(value);
73        if (!count) return false;
74      
75        if (count === 1) {
76            this.map.delete(value);
77        } else {
78            this.map.set(value, count - 1);
79        }
80        this._size--;
81        this.needsSort = true;
82        return true;
83    }
84  
85    has(value: T): boolean {
86        return this.map.has(value);
87    }
88  
89    min(): T | undefined {
90        this.ensureSorted();
91        return this.sorted[0];
92    }
93  
94    max(): T | undefined {
95        this.ensureSorted();
96        return this.sorted[this.sorted.length - 1];
97    }
98  
99    get size(): number {
100        return this._size;
101    }
102  
103    private ensureSorted(): void {
104        if (!this.needsSort) return;
105      
106        this.sorted = [];
107        for (const [value, count] of this.map) {
108            for (let i = 0; i < count; i++) {
109                this.sorted.push(value);
110            }
111        }
112        this.sorted.sort((a, b) => (a as any) - (b as any));
113        this.needsSort = false;
114    }
115}
116

Time and Space Complexity

The time complexity is O(n × log k), where n is the length of the array nums and k is the window size parameter.

Time Complexity Analysis:

The main loop iterates through all n elements in nums
For each iteration:
- l.add(x): O(log k) - insertion into a sorted list of size at most k/2
- l.pop(): O(log k) - removal from a sorted list
- r.add(y): O(log k) - insertion into a sorted list of size at most k/2 + 1
- Potential rebalancing with r.pop(0) and l.add(y): O(log k)
- When i >= k - 1:
  - Removal operations (r.remove() or l.remove()): O(log k)
Since the sliding window has size k, the sorted lists l and r together contain at most k elements
Total: n iterations × O(log k) operations per iteration = O(n × log k)

Space Complexity Analysis: The space complexity is O(k).

The two SortedList objects l and r together store exactly k elements when the window is full
Additional variables (s1, s2, ans, etc.) use O(1) space
Total space: O(k)

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Incorrect Median Selection After Window Sliding

The Pitfall: After removing an element from the sliding window, the balance between left_half and right_half can be disrupted. If we remove an element from right_half when both halves had equal size, left_half will have more elements than right_half. This violates the invariant that the median should always be right_half[0], leading to incorrect cost calculations.

Example Scenario:

Before removal: left_half = [1, 2], right_half = [3, 4] (median = 3)
Remove 3 from right_half: left_half = [1, 2], right_half = [4]
Now left_half has more elements, but we still use right_half[0] = 4 as median (incorrect!)

Solution: Add a rebalancing step after removing elements to ensure right_half always has either equal or one more element than left_half:

# After removing element from the window
if element_to_remove in right_half:
    right_half.remove(element_to_remove)
    right_sum -= element_to_remove
else:
    left_half.remove(element_to_remove)
    left_sum -= element_to_remove

# Rebalance after removal
if len(left_half) > len(right_half):
    largest_left = left_half.pop()
    left_sum -= largest_left
    right_half.add(largest_left)
    right_sum += largest_left

2. Floating Point Precision Issues with Large Numbers

The Pitfall: When dealing with large arrays or large element values, the sums left_sum and right_sum can become very large. This might cause integer overflow in some languages or precision issues when calculating the cost.

Solution: Use appropriate data types and consider modular arithmetic if needed:

# Use integer operations throughout (Python handles big integers well)
# Avoid unnecessary floating-point conversions
cost = (right_sum - median * len(right_half)) + (median * len(left_half) - left_sum)

# Alternative: If working with very large numbers, consider using Decimal for precision
from decimal import Decimal
left_sum = Decimal(0)
right_sum = Decimal(0)

3. Edge Case: k = 1

The Pitfall: When k = 1, every subarray has only one element that's already equal to itself, requiring 0 operations. The two-set approach might overcomplicate this scenario or even fail if not handled properly.

Solution: Add a special case check at the beginning:

def minOperations(self, nums: List[int], k: int) -> int:
    if k == 1:
        return 0
  
    # Continue with the two-set approach for k > 1
    from sortedcontainers import SortedList
    # ... rest of the implementation

Discover Your Strengths and Weaknesses: Take Our 5-Minute Quiz to Tailor Your Study Plan:

Which of the following shows the order of node visit in a Breadth-first Search?

I don't know