Facebook Pixel

1825. Finding MK Average

HardDesignQueueData StreamOrdered SetHeap (Priority Queue)
Leetcode Link

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:

  1. Check stream size: If the stream has fewer than m elements, the MKAverage is -1.

  2. Take last m elements: Otherwise, consider only the last m elements from the stream.

  3. Remove extremes: From these m elements, remove the k smallest values and the k largest values.

  4. 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 parameters m and k.

  • void addElement(int num): Adds a new element num to the stream.

  • int calculateMKAverage(): Returns the MKAverage of the current stream (considering only the last m elements if the stream has at least m elements). Returns -1 if the stream has fewer than m 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)
Quick Interview Experience
Help others by sharing your interview experience
Have you seen this problem before?

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 exactly k elements (or less if we don't have enough total elements)
  • The hi group has exactly k 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:

  1. Queue q: Stores the last m elements in order of insertion, with the head being the oldest element
  2. Three SortedLists:
    • lo: Stores the k smallest elements
    • hi: Stores the k largest elements
    • mid: Stores the remaining middle elements
  3. Variable s: Maintains the running sum of all elements in mid

Implementation Details:

Adding an Element (addElement):

  1. Initial Placement: When a new element num arrives, determine which set it belongs to:

    • If lo is empty or num ≤ max(lo), add to lo
    • Else if hi is empty or num ≥ min(hi), add to hi
    • Otherwise, add to mid and update sum s += num

  2. Maintain Window Size: Add num to queue q. If queue length exceeds m:

    • Remove the oldest element x from the queue
    • Find and remove x from whichever set (lo, mid, or hi) contains it
    • If removing from mid, update sum s -= x

  3. Rebalancing Operations: After insertion/deletion, ensure each set has the correct size:

    • While len(lo) > k: Move max(lo) to mid, update s += max(lo)
    • While len(hi) > k: Move min(hi) to mid, update s += min(hi)
    • While len(lo) < k and mid is not empty: Move min(mid) to lo, update s -= min(mid)
    • While len(hi) < k and mid is not empty: Move max(mid) to hi, update s -= max(mid)

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 lists
  • calculateMKAverage: O(1) since we maintain the sum

Space Complexity:

  • O(m) to store at most m 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 Evaluator

Example 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):

  1. Since lo is empty, add 3 to lo: lo = [3]
  2. Add to queue: q = [3]
  3. Queue size (1) < m (5), no removal needed
  4. Rebalance: len(lo) = 1 = k, no rebalancing needed
  • State: lo = [3], mid = [], hi = [], s = 0

addElement(1):

  1. 1 ≤ max(lo) = 3, so add to lo: lo = [1, 3]
  2. Add to queue: q = [3, 1]
  3. Queue size (2) < m (5), no removal
  4. Rebalance: len(lo) = 2 > k = 1, move max(lo) = 3 to mid
  • State: lo = [1], mid = [3], hi = [], s = 3

addElement(10):

  1. 10 > max(lo) = 1 and hi is empty, add to hi: hi = [10]
  2. Add to queue: q = [3, 1, 10]
  3. Queue size (3) < m (5), no removal
  4. Rebalance: All sets have correct sizes
  • State: lo = [1], mid = [3], hi = [10], s = 3

addElement(5):

  1. 5 > max(lo) = 1 and 5 < min(hi) = 10, add to mid: mid = [3, 5]
  2. Add to queue: q = [3, 1, 10, 5], s = 3 + 5 = 8
  3. Queue size (4) < m (5), no removal
  4. Rebalance: All sets have correct sizes
  • State: lo = [1], mid = [3, 5], hi = [10], s = 8

addElement(7):

  1. 7 > max(lo) = 1 and 7 < min(hi) = 10, add to mid: mid = [3, 5, 7]
  2. Add to queue: q = [3, 1, 10, 5, 7], s = 8 + 7 = 15
  3. Queue size (5) = m (5), no removal yet
  4. 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):

  1. 2 > max(lo) = 1 and 2 < min(hi) = 10, add to mid: mid = [2, 3, 5, 7]
  2. Add to queue: q = [3, 1, 10, 5, 7, 2], s = 15 + 2 = 17
  3. 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]
  4. 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 move O(k) elements, but amortized over many operations, each element is moved a constant number of times, keeping the overall complexity at O(log m)

  • calculateMKAverage(): O(1)

    • Simply returns the pre-calculated sum s divided by (m - 2 * k), which is a constant time operation

Space Complexity: O(m)

  • The deque q stores at most m elements
  • The three SortedLists (lo, mid, hi) together store exactly the same elements as in the deque, totaling m elements
  • Additional variables (s, m, k) use O(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:

  1. Test with edge cases where k = 0 or 2k = m
  2. Verify the sum is correctly maintained by checking sum(middle) == middle_sum during debugging
  3. Test with duplicate values to ensure removal works correctly
  4. Always handle the "excess" case before the "deficit" case in rebalancing
Discover Your Strengths and Weaknesses: Take Our 5-Minute Quiz to Tailor Your Study Plan:

Which data structure is used to implement priority queue?

Recommended Readings

Queue Intro

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

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

Heap Fundamentals

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!

Load More