1825. Finding MK Average
Problem Description
You need to implement a data structure that calculates a special type of average called MKAverage from a stream of integers.
Given two integers m
and k
, the MKAverage is calculated as follows:
-
Check stream size: If the stream has fewer than
m
elements, the MKAverage is-1
. -
Take last m elements: Otherwise, consider only the last
m
elements from the stream. -
Remove extremes: From these
m
elements, remove thek
smallest values and thek
largest values. -
Calculate average: Compute the average of the remaining
m - 2k
elements, rounded down to the nearest integer.
You need to implement a MKAverage
class with three methods:
-
MKAverage(int m, int k)
: Constructor that initializes the object with an empty stream and the parametersm
andk
. -
void addElement(int num)
: Adds a new elementnum
to the stream. -
int calculateMKAverage()
: Returns the MKAverage of the current stream (considering only the lastm
elements if the stream has at leastm
elements). Returns-1
if the stream has fewer thanm
elements.
For example, if m = 5
and k = 1
, and the stream contains [3, 1, 10, 5, 7, 2]
:
- Take the last 5 elements:
[1, 10, 5, 7, 2]
- Remove 1 smallest (
1
) and 1 largest (10
):[5, 7, 2]
- Calculate average:
(5 + 7 + 2) / 3 = 14 / 3 = 4
(rounded down)
Intuition
The key challenge is efficiently maintaining the middle elements (after removing k
smallest and k
largest) as we add new elements and potentially remove old ones from our sliding window of size m
.
The naive approach would be to sort the last m
elements every time we need to calculate the average, but this would be inefficient with O(m log m)
time for each query.
Instead, we can maintain the elements in a way that keeps them already "partitioned" into three groups:
- The
k
smallest elements - The middle
m - 2k
elements - The
k
largest elements
This leads us to think about using three separate ordered data structures. When a new element arrives, we can determine which group it belongs to by comparing it with the boundaries of these groups (the maximum of the smallest group and the minimum of the largest group).
Since we only care about the last m
elements, we need a queue to track the order of insertion. When the queue exceeds size m
, we remove the oldest element from whichever group contains it.
The crucial insight is that after adding or removing elements, we might need to rebalance the groups to ensure:
- The
lo
group has exactlyk
elements (or less if we don't have enough total elements) - The
hi
group has exactlyk
elements (or less if we don't have enough total elements) - The
mid
group contains everything else
This rebalancing involves moving elements between adjacent groups. For example, if lo
has more than k
elements, we move its largest element to mid
. If lo
has fewer than k
elements and mid
is not empty, we move the smallest element from mid
to lo
.
By maintaining the sum of only the middle group elements, we can calculate the MKAverage in O(1)
time by simply dividing this sum by m - 2k
.
Learn more about Queue, Data Stream and Heap (Priority Queue) patterns.
Solution Approach
The solution uses three ordered sets and a queue to efficiently maintain the required data structure:
Data Structures:
- Queue
q
: Stores the lastm
elements in order of insertion, with the head being the oldest element - Three SortedLists:
lo
: Stores thek
smallest elementshi
: Stores thek
largest elementsmid
: Stores the remaining middle elements
- Variable
s
: Maintains the running sum of all elements inmid
Implementation Details:
Adding an Element (addElement
):
-
Initial Placement: When a new element
num
arrives, determine which set it belongs to:- If
lo
is empty ornum ≤ max(lo)
, add tolo
- Else if
hi
is empty ornum ≥ min(hi)
, add tohi
- Otherwise, add to
mid
and update sums += num
- If
-
Maintain Window Size: Add
num
to queueq
. If queue length exceedsm
:- Remove the oldest element
x
from the queue - Find and remove
x
from whichever set (lo
,mid
, orhi
) contains it - If removing from
mid
, update sums -= x
- Remove the oldest element
-
Rebalancing Operations: After insertion/deletion, ensure each set has the correct size:
- While
len(lo) > k
: Movemax(lo)
tomid
, updates += max(lo)
- While
len(hi) > k
: Movemin(hi)
tomid
, updates += min(hi)
- While
len(lo) < k
andmid
is not empty: Movemin(mid)
tolo
, updates -= min(mid)
- While
len(hi) < k
andmid
is not empty: Movemax(mid)
tohi
, updates -= max(mid)
- While
Calculating MKAverage (calculateMKAverage
):
- If queue length <
m
, return-1
- Otherwise, return
s // (m - 2*k)
(integer division for rounding down)
Time Complexity:
addElement
:O(log m)
for insertions/removals in sorted listscalculateMKAverage
:O(1)
since we maintain the sum
Space Complexity:
O(m)
to store at mostm
elements across all data structures
The key insight is that by maintaining these three partitioned sets and rebalancing them after each operation, we avoid having to sort the entire window repeatedly and can calculate the average in constant time.
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 walk through an example with m = 5
, k = 1
, and the following sequence of operations:
Initial State:
lo = []
,mid = []
,hi = []
,q = []
,s = 0
addElement(3):
- Since
lo
is empty, add 3 tolo
:lo = [3]
- Add to queue:
q = [3]
- Queue size (1) < m (5), no removal needed
- Rebalance:
len(lo) = 1 = k
, no rebalancing needed
- State:
lo = [3]
,mid = []
,hi = []
,s = 0
addElement(1):
- 1 ≤ max(lo) = 3, so add to
lo
:lo = [1, 3]
- Add to queue:
q = [3, 1]
- Queue size (2) < m (5), no removal
- Rebalance:
len(lo) = 2 > k = 1
, move max(lo) = 3 tomid
- State:
lo = [1]
,mid = [3]
,hi = []
,s = 3
addElement(10):
- 10 > max(lo) = 1 and
hi
is empty, add tohi
:hi = [10]
- Add to queue:
q = [3, 1, 10]
- Queue size (3) < m (5), no removal
- Rebalance: All sets have correct sizes
- State:
lo = [1]
,mid = [3]
,hi = [10]
,s = 3
addElement(5):
- 5 > max(lo) = 1 and 5 < min(hi) = 10, add to
mid
:mid = [3, 5]
- Add to queue:
q = [3, 1, 10, 5]
,s = 3 + 5 = 8
- Queue size (4) < m (5), no removal
- Rebalance: All sets have correct sizes
- State:
lo = [1]
,mid = [3, 5]
,hi = [10]
,s = 8
addElement(7):
- 7 > max(lo) = 1 and 7 < min(hi) = 10, add to
mid
:mid = [3, 5, 7]
- Add to queue:
q = [3, 1, 10, 5, 7]
,s = 8 + 7 = 15
- Queue size (5) = m (5), no removal yet
- Rebalance: All sets have correct sizes
- State:
lo = [1]
,mid = [3, 5, 7]
,hi = [10]
,s = 15
calculateMKAverage():
- Queue size = 5 ≥ m = 5
- Return
s // (m - 2*k) = 15 // 3 = 5
addElement(2):
- 2 > max(lo) = 1 and 2 < min(hi) = 10, add to
mid
:mid = [2, 3, 5, 7]
- Add to queue:
q = [3, 1, 10, 5, 7, 2]
,s = 15 + 2 = 17
- Queue size (6) > m (5), remove oldest element 3:
- Find 3 in
mid
, remove it:mid = [2, 5, 7]
- Update sum:
s = 17 - 3 = 14
- Remove from queue:
q = [1, 10, 5, 7, 2]
- Find 3 in
- Rebalance: All sets have correct sizes
- Final state:
lo = [1]
,mid = [2, 5, 7]
,hi = [10]
,s = 14
calculateMKAverage():
- The last 5 elements are
[1, 10, 5, 7, 2]
- After removing k=1 smallest (1) and k=1 largest (10):
[2, 5, 7]
- Return
s // (m - 2*k) = 14 // 3 = 4
This walkthrough demonstrates how the three sets maintain the partitioning, how rebalancing ensures correct sizes, and how the running sum enables O(1) average calculation.
Solution Implementation
1from collections import deque
2from sortedcontainers import SortedList
3
4class MKAverage:
5 def __init__(self, m: int, k: int):
6 """
7 Initialize the MKAverage data structure.
8
9 Args:
10 m: The number of most recent elements to consider
11 k: The number of smallest and largest elements to exclude
12 """
13 self.m = m # Size of sliding window
14 self.k = k # Number of elements to exclude from each end
15 self.middle_sum = 0 # Sum of middle elements (excluding k smallest and k largest)
16 self.queue = deque() # Queue to maintain last m elements in order
17
18 # Three sorted lists to maintain k smallest, middle, and k largest elements
19 self.smallest = SortedList() # Stores k smallest elements
20 self.middle = SortedList() # Stores middle elements (m - 2k elements)
21 self.largest = SortedList() # Stores k largest elements
22
23 def addElement(self, num: int) -> None:
24 """
25 Add a new element to the data structure.
26
27 Args:
28 num: The number to add
29 """
30 # Determine which sorted list the new element belongs to
31 if not self.smallest or num <= self.smallest[-1]:
32 # Number belongs to smallest k elements
33 self.smallest.add(num)
34 elif not self.largest or num >= self.largest[0]:
35 # Number belongs to largest k elements
36 self.largest.add(num)
37 else:
38 # Number belongs to middle elements
39 self.middle.add(num)
40 self.middle_sum += num
41
42 # Add element to queue
43 self.queue.append(num)
44
45 # Remove oldest element if queue exceeds m elements
46 if len(self.queue) > self.m:
47 oldest = self.queue.popleft()
48
49 # Remove oldest element from appropriate sorted list
50 if oldest in self.smallest:
51 self.smallest.remove(oldest)
52 elif oldest in self.largest:
53 self.largest.remove(oldest)
54 else:
55 self.middle.remove(oldest)
56 self.middle_sum -= oldest
57
58 # Rebalance the three sorted lists to maintain invariants
59
60 # Move excess elements from smallest to middle
61 while len(self.smallest) > self.k:
62 element = self.smallest.pop() # Remove largest from smallest
63 self.middle.add(element)
64 self.middle_sum += element
65
66 # Move excess elements from largest to middle
67 while len(self.largest) > self.k:
68 element = self.largest.pop(0) # Remove smallest from largest
69 self.middle.add(element)
70 self.middle_sum += element
71
72 # Move elements from middle to smallest if needed
73 while len(self.smallest) < self.k and self.middle:
74 element = self.middle.pop(0) # Remove smallest from middle
75 self.smallest.add(element)
76 self.middle_sum -= element
77
78 # Move elements from middle to largest if needed
79 while len(self.largest) < self.k and self.middle:
80 element = self.middle.pop() # Remove largest from middle
81 self.largest.add(element)
82 self.middle_sum -= element
83
84 def calculateMKAverage(self) -> int:
85 """
86 Calculate the MKAverage of current elements.
87
88 Returns:
89 The average of middle elements (excluding k smallest and k largest),
90 or -1 if there are fewer than m elements
91 """
92 # Return -1 if we don't have enough elements
93 if len(self.queue) < self.m:
94 return -1
95
96 # Calculate average of middle elements (integer division)
97 middle_count = self.m - 2 * self.k
98 return self.middle_sum // middle_count
99
100
101# Your MKAverage object will be instantiated and called as such:
102# obj = MKAverage(m, k)
103# obj.addElement(num)
104# param_2 = obj.calculateMKAverage()
105
1class MKAverage {
2 // Configuration parameters
3 private int m; // Size of the sliding window
4 private int k; // Number of smallest/largest elements to exclude
5
6 // State variables
7 private long middleSum; // Sum of elements in the middle TreeMap
8 private int lowSize; // Current size of low TreeMap
9 private int highSize; // Current size of high TreeMap
10
11 // Data structures
12 private Deque<Integer> queue = new ArrayDeque<>(); // Sliding window queue
13 private TreeMap<Integer, Integer> low = new TreeMap<>(); // Stores k smallest elements
14 private TreeMap<Integer, Integer> middle = new TreeMap<>(); // Stores middle elements
15 private TreeMap<Integer, Integer> high = new TreeMap<>(); // Stores k largest elements
16
17 /**
18 * Initializes the MKAverage object with window size m and exclusion count k
19 * @param m The size of the sliding window
20 * @param k The number of smallest and largest elements to exclude
21 */
22 public MKAverage(int m, int k) {
23 this.m = m;
24 this.k = k;
25 }
26
27 /**
28 * Adds a new element to the data stream
29 * @param num The element to add
30 */
31 public void addElement(int num) {
32 // Step 1: Add the new element to appropriate TreeMap
33 if (low.isEmpty() || num <= low.lastKey()) {
34 // Element belongs in low TreeMap (smallest k elements)
35 low.merge(num, 1, Integer::sum);
36 lowSize++;
37 } else if (high.isEmpty() || num >= high.firstKey()) {
38 // Element belongs in high TreeMap (largest k elements)
39 high.merge(num, 1, Integer::sum);
40 highSize++;
41 } else {
42 // Element belongs in middle TreeMap
43 middle.merge(num, 1, Integer::sum);
44 middleSum += num;
45 }
46
47 // Step 2: Add element to sliding window queue
48 queue.offer(num);
49
50 // Step 3: Remove oldest element if window size exceeds m
51 if (queue.size() > m) {
52 int removedElement = queue.poll();
53
54 // Remove from appropriate TreeMap
55 if (low.containsKey(removedElement)) {
56 if (low.merge(removedElement, -1, Integer::sum) == 0) {
57 low.remove(removedElement);
58 }
59 lowSize--;
60 } else if (high.containsKey(removedElement)) {
61 if (high.merge(removedElement, -1, Integer::sum) == 0) {
62 high.remove(removedElement);
63 }
64 highSize--;
65 } else {
66 if (middle.merge(removedElement, -1, Integer::sum) == 0) {
67 middle.remove(removedElement);
68 }
69 middleSum -= removedElement;
70 }
71 }
72
73 // Step 4: Rebalance TreeMaps to maintain size constraints
74
75 // Move excess elements from low to middle
76 while (lowSize > k) {
77 int maxLowElement = low.lastKey();
78 if (low.merge(maxLowElement, -1, Integer::sum) == 0) {
79 low.remove(maxLowElement);
80 }
81 middle.merge(maxLowElement, 1, Integer::sum);
82 middleSum += maxLowElement;
83 lowSize--;
84 }
85
86 // Move excess elements from high to middle
87 while (highSize > k) {
88 int minHighElement = high.firstKey();
89 if (high.merge(minHighElement, -1, Integer::sum) == 0) {
90 high.remove(minHighElement);
91 }
92 middle.merge(minHighElement, 1, Integer::sum);
93 middleSum += minHighElement;
94 highSize--;
95 }
96
97 // Move elements from middle to low if needed
98 while (lowSize < k && !middle.isEmpty()) {
99 int minMiddleElement = middle.firstKey();
100 if (middle.merge(minMiddleElement, -1, Integer::sum) == 0) {
101 middle.remove(minMiddleElement);
102 }
103 middleSum -= minMiddleElement;
104 low.merge(minMiddleElement, 1, Integer::sum);
105 lowSize++;
106 }
107
108 // Move elements from middle to high if needed
109 while (highSize < k && !middle.isEmpty()) {
110 int maxMiddleElement = middle.lastKey();
111 if (middle.merge(maxMiddleElement, -1, Integer::sum) == 0) {
112 middle.remove(maxMiddleElement);
113 }
114 middleSum -= maxMiddleElement;
115 high.merge(maxMiddleElement, 1, Integer::sum);
116 highSize++;
117 }
118 }
119
120 /**
121 * Calculates the MKAverage of the current stream
122 * @return The average of middle elements, or -1 if not enough elements
123 */
124 public int calculateMKAverage() {
125 // Check if we have enough elements (at least m elements)
126 if (queue.size() < m) {
127 return -1;
128 }
129
130 // Calculate average of middle elements
131 // Total middle elements = queue.size() - 2*k
132 return (int) (middleSum / (queue.size() - k * 2));
133 }
134}
135
136/**
137 * Your MKAverage object will be instantiated and called as such:
138 * MKAverage obj = new MKAverage(m, k);
139 * obj.addElement(num);
140 * int param_2 = obj.calculateMKAverage();
141 */
142
1class MKAverage {
2private:
3 int m; // Size of the sliding window
4 int k; // Number of smallest and largest elements to exclude
5 long long midSum; // Sum of elements in the middle multiset
6 queue<int> window; // Queue to maintain the sliding window of last m elements
7 multiset<int> smallest; // Multiset containing k smallest elements
8 multiset<int> middle; // Multiset containing middle elements (excluding k smallest and k largest)
9 multiset<int> largest; // Multiset containing k largest elements
10
11public:
12 MKAverage(int m, int k) {
13 this->m = m;
14 this->k = k;
15 this->midSum = 0;
16 }
17
18 void addElement(int num) {
19 // Add new element to appropriate multiset based on current boundaries
20 if (smallest.empty() || num <= *smallest.rbegin()) {
21 smallest.insert(num);
22 } else if (largest.empty() || num >= *largest.begin()) {
23 largest.insert(num);
24 } else {
25 middle.insert(num);
26 midSum += num;
27 }
28
29 // Add element to sliding window
30 window.push(num);
31
32 // Remove oldest element if window size exceeds m
33 if (window.size() > m) {
34 int oldestElement = window.front();
35 window.pop();
36
37 // Remove the oldest element from appropriate multiset
38 if (smallest.find(oldestElement) != smallest.end()) {
39 smallest.erase(smallest.find(oldestElement));
40 } else if (largest.find(oldestElement) != largest.end()) {
41 largest.erase(largest.find(oldestElement));
42 } else {
43 middle.erase(middle.find(oldestElement));
44 midSum -= oldestElement;
45 }
46 }
47
48 // Rebalance multisets to maintain size constraints
49
50 // Move excess elements from smallest to middle
51 while (smallest.size() > k) {
52 int maxInSmallest = *smallest.rbegin();
53 smallest.erase(prev(smallest.end()));
54 middle.insert(maxInSmallest);
55 midSum += maxInSmallest;
56 }
57
58 // Move excess elements from largest to middle
59 while (largest.size() > k) {
60 int minInLargest = *largest.begin();
61 largest.erase(largest.begin());
62 middle.insert(minInLargest);
63 midSum += minInLargest;
64 }
65
66 // Fill smallest multiset from middle if needed
67 while (smallest.size() < k && !middle.empty()) {
68 int minInMiddle = *middle.begin();
69 middle.erase(middle.begin());
70 midSum -= minInMiddle;
71 smallest.insert(minInMiddle);
72 }
73
74 // Fill largest multiset from middle if needed
75 while (largest.size() < k && !middle.empty()) {
76 int maxInMiddle = *middle.rbegin();
77 middle.erase(prev(middle.end()));
78 midSum -= maxInMiddle;
79 largest.insert(maxInMiddle);
80 }
81 }
82
83 int calculateMKAverage() {
84 // Return -1 if we don't have enough elements
85 if (window.size() < m) {
86 return -1;
87 }
88
89 // Calculate average of middle elements
90 // Total elements - 2k (k smallest + k largest)
91 return midSum / (window.size() - 2 * k);
92 }
93};
94
95/**
96 * Your MKAverage object will be instantiated and called as such:
97 * MKAverage* obj = new MKAverage(m, k);
98 * obj->addElement(num);
99 * int param_2 = obj->calculateMKAverage();
100 */
101
1// Global variables for MKAverage implementation
2let m: number; // Size of the sliding window
3let k: number; // Number of smallest and largest elements to exclude
4let midSum: number; // Sum of elements in the middle set
5let window: number[]; // Array to maintain the sliding window of last m elements
6let smallest: number[]; // Array containing k smallest elements (kept sorted)
7let middle: number[]; // Array containing middle elements (excluding k smallest and k largest, kept sorted)
8let largest: number[]; // Array containing k largest elements (kept sorted)
9
10/**
11 * Initializes the MKAverage data structure with given parameters
12 * @param mValue - Size of the sliding window
13 * @param kValue - Number of smallest and largest elements to exclude
14 */
15function MKAverage(mValue: number, kValue: number): void {
16 m = mValue;
17 k = kValue;
18 midSum = 0;
19 window = [];
20 smallest = [];
21 middle = [];
22 largest = [];
23}
24
25/**
26 * Binary search to find insertion position in a sorted array
27 * @param arr - Sorted array to search in
28 * @param target - Value to find position for
29 * @returns Index where target should be inserted
30 */
31function binarySearchInsertPosition(arr: number[], target: number): number {
32 let left = 0;
33 let right = arr.length;
34
35 while (left < right) {
36 const mid = Math.floor((left + right) / 2);
37 if (arr[mid] < target) {
38 left = mid + 1;
39 } else {
40 right = mid;
41 }
42 }
43
44 return left;
45}
46
47/**
48 * Inserts a value into a sorted array maintaining sort order
49 * @param arr - Sorted array to insert into
50 * @param value - Value to insert
51 */
52function insertSorted(arr: number[], value: number): void {
53 const pos = binarySearchInsertPosition(arr, value);
54 arr.splice(pos, 0, value);
55}
56
57/**
58 * Removes first occurrence of a value from a sorted array
59 * @param arr - Sorted array to remove from
60 * @param value - Value to remove
61 */
62function removeSorted(arr: number[], value: number): void {
63 const pos = binarySearchInsertPosition(arr, value);
64 if (pos < arr.length && arr[pos] === value) {
65 arr.splice(pos, 1);
66 }
67}
68
69/**
70 * Adds a new element to the data structure
71 * @param num - The element to add
72 */
73function addElement(num: number): void {
74 // Add new element to appropriate set based on current boundaries
75 if (smallest.length === 0 || num <= smallest[smallest.length - 1]) {
76 insertSorted(smallest, num);
77 } else if (largest.length === 0 || num >= largest[0]) {
78 insertSorted(largest, num);
79 } else {
80 insertSorted(middle, num);
81 midSum += num;
82 }
83
84 // Add element to sliding window
85 window.push(num);
86
87 // Remove oldest element if window size exceeds m
88 if (window.length > m) {
89 const oldestElement = window.shift()!;
90
91 // Remove the oldest element from appropriate set
92 if (smallest.includes(oldestElement)) {
93 removeSorted(smallest, oldestElement);
94 } else if (largest.includes(oldestElement)) {
95 removeSorted(largest, oldestElement);
96 } else {
97 removeSorted(middle, oldestElement);
98 midSum -= oldestElement;
99 }
100 }
101
102 // Rebalance sets to maintain size constraints
103
104 // Move excess elements from smallest to middle
105 while (smallest.length > k) {
106 const maxInSmallest = smallest.pop()!;
107 insertSorted(middle, maxInSmallest);
108 midSum += maxInSmallest;
109 }
110
111 // Move excess elements from largest to middle
112 while (largest.length > k) {
113 const minInLargest = largest.shift()!;
114 insertSorted(middle, minInLargest);
115 midSum += minInLargest;
116 }
117
118 // Fill smallest set from middle if needed
119 while (smallest.length < k && middle.length > 0) {
120 const minInMiddle = middle.shift()!;
121 midSum -= minInMiddle;
122 insertSorted(smallest, minInMiddle);
123 }
124
125 // Fill largest set from middle if needed
126 while (largest.length < k && middle.length > 0) {
127 const maxInMiddle = middle.pop()!;
128 midSum -= maxInMiddle;
129 insertSorted(largest, maxInMiddle);
130 }
131}
132
133/**
134 * Calculates the MK average of the current window
135 * @returns The MK average, or -1 if not enough elements
136 */
137function calculateMKAverage(): number {
138 // Return -1 if we don't have enough elements
139 if (window.length < m) {
140 return -1;
141 }
142
143 // Calculate average of middle elements
144 // Total elements - 2k (k smallest + k largest)
145 return Math.floor(midSum / (window.length - 2 * k));
146}
147
Time and Space Complexity
Time Complexity:
-
addElement(num)
:O(log m)
per operation- Adding to
SortedList
(lo, mid, or hi):O(log m)
- Appending to deque:
O(1)
- When queue exceeds size
m
, removing from deque:O(1)
and removing from one of the SortedLists:O(log m)
- Rebalancing operations (while loops): Each element movement involves removing from one SortedList (
O(log m)
) and adding to another (O(log m)
). In the worst case, we might need to moveO(k)
elements, but amortized over many operations, each element is moved a constant number of times, keeping the overall complexity atO(log m)
- Adding to
-
calculateMKAverage()
:O(1)
- Simply returns the pre-calculated sum
s
divided by(m - 2 * k)
, which is a constant time operation
- Simply returns the pre-calculated sum
Space Complexity: O(m)
- The deque
q
stores at mostm
elements - The three SortedLists (
lo
,mid
,hi
) together store exactly the same elements as in the deque, totalingm
elements - Additional variables (
s
,m
,k
) useO(1)
space - Total space complexity is
O(m)
Learn more about how to find time and space complexity quickly.
Common Pitfalls
1. Incorrect Rebalancing Order
One of the most critical pitfalls is performing the rebalancing operations in the wrong order. The rebalancing must follow a specific sequence to maintain the invariants correctly.
Wrong Approach:
# Moving elements to smallest/largest before handling excess
while len(self.smallest) < self.k and self.middle:
element = self.middle.pop(0)
self.smallest.add(element)
self.middle_sum -= element
while len(self.smallest) > self.k:
element = self.smallest.pop()
self.middle.add(element)
self.middle_sum += element
Why it's wrong: If you try to fill smallest
or largest
before removing excess elements, you might move elements unnecessarily and corrupt the data structure state.
Correct Approach: Always handle excess elements first (moving from smallest
/largest
to middle
), then handle deficits (moving from middle
to smallest
/largest
).
2. Incorrect Element Removal from SortedList
When removing the oldest element from the queue, using the wrong removal method can lead to removing the wrong element or errors.
Wrong Approach:
if oldest in self.smallest: self.smallest.pop(0) # Wrong! This removes the minimum, not 'oldest'
Correct Approach:
if oldest in self.smallest: self.smallest.remove(oldest) # Removes the specific element
3. Not Handling Edge Cases in Initial Placement
When adding a new element, comparing with empty lists can cause errors.
Wrong Approach:
if num <= self.smallest[-1]: # Error if smallest is empty! self.smallest.add(num)
Correct Approach:
if not self.smallest or num <= self.smallest[-1]: self.smallest.add(num)
4. Forgetting to Update the Sum
The middle_sum
must be updated whenever elements move in or out of the middle
list.
Wrong Approach:
# Forgetting to update sum when removing from middle if oldest in self.middle: self.middle.remove(oldest) # Missing: self.middle_sum -= oldest
Correct Approach: Always update middle_sum
when:
- Adding to
middle
:self.middle_sum += num
- Removing from
middle
:self.middle_sum -= num
- Moving elements in/out during rebalancing
5. Using Regular List Instead of SortedList
Using Python's regular list and manually sorting will result in O(m log m) complexity for each operation.
Wrong Approach:
self.smallest = [] # Regular list # Later... self.smallest.append(num) self.smallest.sort() # O(m log m) every time!
Correct Approach: Use SortedList
from sortedcontainers
which maintains sorted order with O(log m) insertions and removals.
6. Integer Division vs Float Division
The problem requires rounding down to the nearest integer.
Wrong Approach:
return self.middle_sum / middle_count # Returns float
Correct Approach:
return self.middle_sum // middle_count # Integer division (rounds down)
Prevention Tips:
- Test with edge cases where
k = 0
or2k = m
- Verify the sum is correctly maintained by checking
sum(middle) == middle_sum
during debugging - Test with duplicate values to ensure removal works correctly
- Always handle the "excess" case before the "deficit" case in rebalancing
Which data structure is used to implement recursion?
Recommended Readings
Queue Intro Think of the last time you stood in line to buy movie tickets The first person in line gets their ticket first and any newcomers join the end of the line This system where the first person to arrive is the first to be served is a queue in real
Median of Data Stream Given a stream of numbers find the median number at any given time accurate to 1 decimal place Example add_number 1 add_number 2 add_number 3 get_median 2 0 add_number 4 get_median 2 5 Try it yourself Explanation Intuition Brute force way is to sort the entire list every time we get a new number This would be O N log N for each add_number operation One step up is to notice that the list is sorted before we add any new number to it So every
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!