3422. Minimum Operations to Make Subarray Elements Equal 🔒
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.
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 tol
) - 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 windowr
: SortedList storing the larger half of elements (the minimum ofr
is the median)s1
: Sum of all elements inl
s2
: Sum of all elements inr
Algorithm Steps:
-
Building the Initial Window:
- For each element
x
innums
, we first add it to the left setl
and updates1
- We then pop the maximum element from
l
and move it tor
to maintain balance - This ensures elements are properly distributed between the two sets
- For each element
-
Maintaining Balance:
- After adding to
r
, we check iflen(r) - len(l) > 1
- If true, we move the minimum element from
r
back tol
- This maintains the invariant that
r
has at most one more element thanl
- After adding to
-
Calculating Operations for Each Window:
- Once we have
k
elements (wheni >= 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 inl
up to the medians2 - r[0] * len(r)
: Operations to move all elements inr
down to the median
- Once we have
-
Sliding the Window:
- Remove the leftmost element of the previous window (
nums[j]
wherej = 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
- Remove the leftmost element of the previous window (
-
Tracking the Minimum:
- We maintain
ans
to track the minimum operations across all windows - Return
ans
after processing all possible subarrays of sizek
- We maintain
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 EvaluatorExample 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) tor
: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) tor
: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) tol
: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) tor
: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
- Window is
- Remove leftmost element (nums[0] = 1) from window:
- 1 is in
l
, so remove it:l = []
,s1 = 0
- 1 is in
Step 4: Process nums[3] = 6
- Add 6 to
l
:l = [6]
,s1 = 6
- Move max of
l
(which is 6) tor
: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) tol
: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
- Window is
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 innums
- For each iteration:
l.add(x)
:O(log k)
- insertion into a sorted list of size at mostk/2
l.pop()
:O(log k)
- removal from a sorted listr.add(y)
:O(log k)
- insertion into a sorted list of size at mostk/2 + 1
- Potential rebalancing with
r.pop(0)
andl.add(y)
:O(log k)
- When
i >= k - 1
:- Removal operations (
r.remove()
orl.remove()
):O(log k)
- Removal operations (
- Since the sliding window has size
k
, the sorted listsl
andr
together contain at mostk
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
objectsl
andr
together store exactlyk
elements when the window is full - Additional variables (
s1
,s2
,ans
, etc.) useO(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 useright_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
Which algorithm should you use to find a node that is close to the root of the tree?
Recommended Readings
Math for Technical Interviews How much math do I need to know for technical interviews The short answer is about high school level math Computer science is often associated with math and some universities even place their computer science department under the math faculty However the reality is that you
https assets algo monster cover_photos stack svg Sliding Window Maximum Monotonic Stack We have an array and a sliding window defined by a start index and an end index The sliding window moves from left of the array to right There are always k elements in the window The window
https assets algo monster cover_photos heap svg Priority Queue and Heap What is the relationship between priority queue and heap Priority Queue is an Abstract Data Type and Heap is the concrete data structure we use to implement a priority queue Priority Queue A priority queue is a data structure
Want a Structured Path to Master System Design Too? Don’t Miss This!