The problem
The median of a list of numbers is the middle one once they are sorted; with an even count, it is the average of the two middle ones — the median of 2, 3, 4 is 3, of 2, 3 is 2.5.
Design addNum(num), which adds a number, and findMedian(), which returns the median of all numbers added so far.
Examples
- Input
["MedianFinder", "addNum", "addNum", "findMedian", "addNum", "findMedian"] [[], [1], [2], [], [3], []]
- Output
[null, null, null, 1.5, null, 2]
Constraints
- −10⁵ ≤ num ≤ 10⁵
findMedianis only called after at least one number is added.- At most 5 × 10⁴ calls in total.
The idea
The median only needs the one or two numbers in the middle. Split the numbers into a lower half and an upper half: keep the lower half in a max-heap (its largest on top) and the upper half in a min-heap (its smallest on top). The two tops are the middle.
Keep the halves the same size, or the lower one bigger by one. To add a number, push it into the lower half, then move the lower half’s largest into the upper half — so everything below stays below everything above — and if the upper half is now bigger, move its smallest back. The median is the lower top, or the average of both tops.
- Time
- O(log n) to add, O(1) to find
- Space
- O(n)
Solution · every language run against every case
class MedianFinder: def __init__(self): # low: the smaller half, as a max-heap (values negated); high: the larger half, a min-heap. # low holds the same number as high, or one more. self.low, self.high = [], [] def addNum(self, num: int) -> None: heappush(self.low, -num) heappush(self.high, -heappop(self.low)) # the largest of the low half moves up if len(self.high) > len(self.low): heappush(self.low, -heappop(self.high)) # rebalance def findMedian(self) -> float: if len(self.low) > len(self.high): return float(-self.low[0]) return (-self.low[0] + self.high[0]) / 2