The main idea is to use a max-heap to store the smaller half of the numbers and a min-heap to store the larger half. The max-heap allows you to efficiently fetch the largest element of the smaller half, and the min-heap allows you to efficiently fetch the smallest element of the larger half. To maintain the balance, ensure that the number of elements in the max-heap is either equal to or one more than the number of elements in the min-heap.
The time complexity for adding a number is O(log n) due to heap adjustments. Finding the median has a time complexity of O(1). Space complexity is O(n), where n is the number of elements.
1import heapq
2
3class MedianFinder:
4 def __init__(self):
5 self.small = [] # max-heap (inverted to min-heap by storing negative values)
6 self.large = [] # min-heap
7
8 def addNum(self, num: int) -> None:
9 heapq.heappush(self.small, -num)
10 if self.small and self.large and (-self.small[0] > self.large[0]):
11 val = -heapq.heappop(self.small)
12 heapq.heappush(self.large, val)
13 if len(self.small) > len(self.large) + 1:
14 val = -heapq.heappop(self.small)
15 heapq.heappush(self.large, val)
16 if len(self.large) > len(self.small):
17 val = heapq.heappop(self.large)
18 heapq.heappush(self.small, -val)
19
20 def findMedian(self) -> float:
21 if len(self.small) > len(self.large):
22 return -self.small[0]
23 return (-self.small[0] + self.large[0]) / 2
The solution uses two heaps (small and large). 'small' is a max-heap that stores the smaller half of the numbers, and 'large' is a min-heap that stores the larger half. Each time a number is added, it is pushed to the max-heap first. If the max-heap's largest is greater than the min-heap's smallest, balance them by moving numbers between heaps. The median can be derived based on the sizes of the heaps.
An alternative approach is to use a balanced binary search tree (BST) to keep track of the numbers. This allows maintaining a sorted order of elements, which can make finding the median efficient since a BST can give average-case time complexity of O(log n) for insertion and O(1) for median finding.
Insertion and balancing in the sorted set take O(log n) time, and findings of the median are O(1). Space complexity is O(n).
1using System;
2using System.Collections.Generic;
3
4public class MedianFinder {
5 private SortedSet<(int, int)> small;
6 private SortedSet<(int, int)> large;
7 private int counter = 0;
8
9 public MedianFinder() {
10 small = new SortedSet<(int, int)>();
11 large = new SortedSet<(int, int)>();
12 }
13
14 public void AddNum(int num) {
15 small.Add((num, counter++));
16 large.Add(small.Max);
17 small.Remove(small.Max);
18 if (large.Count > small.Count) {
19 small.Add(large.Min);
20 large.Remove(large.Min);
21 }
22 }
23
24 public double FindMedian() {
25 if (small.Count > large.Count) {
26 return small.Max.Item1;
27 }
28 return (small.Max.Item1 + large.Min.Item1) / 2.0;
29 }
30}
This C# implementation uses two sorted sets as a stand-in for a balanced BST. Every number is inserted into the smaller set first, followed by balancing between the two sets. The median is retrieved based on the sizes of the sets and their boundary values.