Find Median from Data Stream
ID: 81; hard
Solution 1 (Java)
public class Solution {
PriorityQueue<Integer> minHeap, maxHeap;
int median;
boolean isFirstNum;
public Solution() {
minHeap = new PriorityQueue<>();
maxHeap = new PriorityQueue<>(new Comparator<Integer>() {
public int compare(Integer x, Integer y) {
return y - x;
}
});
isFirstNum = true;
}
/**
* @param val: An integer
* @return: nothing
*/
public void add(int val) {
if (isFirstNum) {
median = val;
isFirstNum = false;
return;
}
if (val < median) {
maxHeap.offer(val);
} else {
minHeap.offer(val);
}
if (maxHeap.size() > minHeap.size()) {
minHeap.offer(median);
median = maxHeap.poll();
}
if (maxHeap.size() < minHeap.size() - 1) {
maxHeap.offer(median);
median = minHeap.poll();
}
}
/**
* @return: return the median of the data stream
*/
public int getMedian() {
return median;
}
}
Notes
We used a
maxHeap
and aminHeap
to maintain the data stream. ThemaxHeap
stores numbers that are less than themedian
. TheminHeap
stores numbers that are larger than themedian
. Then, we update the median whenever the two heaps have sizes that differ by 2 (to be more precise, see the condition in code). By construction, the results by polling either theminHeap
or themaxHeap
will be the currentmedian
. Eventually, we just return the currentmedian
.The reason is that we need to find the median of
k
numbers, so it is natural to find thek/2
th largest number and thek/2
th smallest number (first number in each heap).
Last updated
Was this helpful?