Find Median from Data Stream
Median is the middle value in an ordered integer list. If the size of the list is even, there are two middle values and the median is their average. For example, for arr = [1, 2, 3, 4], the median is (2 + 3) / 2 = 2.5. You are given an array of integers nums and an integer k. You need to find the kth largest element in the array.
Constraints:
- 1 <= nums.length <= 10^5
- 1 <= k <= nums.length
Examples:
Input: [1, 2] and then [1, 3] and then [2, 3]
Output: [1.5, 2, 2]
Explanation: After adding 1 and 2, the median is 1.5. After adding 1 and 3, the median is 2. After adding 2 and 3, the median is 2.
Solutions
Using Two Heaps
We use two heaps to solve this problem. The max heap stores the smaller half of the numbers and the min heap stores the larger half. We ensure that the max heap always has one more element than the min heap. When we add a new number, we first add it to the max heap. Then we remove the largest number from the max heap and add it to the min heap. If the max heap has less elements than the min heap, we remove the smallest number from the min heap and add it to the max heap. The median is then the average of the largest number in the max heap and the smallest number in the min heap if the total number of elements is even, or the largest number in the max heap if the total number of elements is odd.
import java.util.PriorityQueue;
class MedianFinder {
PriorityQueue<Integer> maxHeap;
PriorityQueue<Integer> minHeap;
public MedianFinder() {
maxHeap = new PriorityQueue<>((a, b) -> b - a);
minHeap = new PriorityQueue<>();
}
public void addNum(int num) {
maxHeap.offer(num);
minHeap.offer(maxHeap.poll());
if (maxHeap.size() < minHeap.size()) {
maxHeap.offer(minHeap.poll());
}
}
public double findMedian() {
if (maxHeap.size() == minHeap.size()) {
return (maxHeap.peek() + minHeap.peek()) / 2.0;
}
else {
return maxHeap.peek();
}
}
}Follow-up:
How would you optimize the solution if the input stream is very large?
