Find Median from Data Stream
Question
Median is the middle value in an ordered integer list. If the size of the list is even, there is no middle value. So the median is the mean of the two middle value.
Examples: [2,3,4] , the median is 3
[2,3], the median is (2 + 3) / 2 = 2.5
Design a data structure that supports the following two operations:
void addNum(int num) - Add a integer number from the data stream to the data structure. double findMedian() - Return the median of all elements so far.
Solution
這題是看了九章之後來寫的,九章的教法是,中位數的右邊用 minQueue;左邊用 maxQueue,然後維護這兩個 queue 的大小,進而求出中位數。所以我自己在實作的時候,中間用了個 deque,因為有時候得從 minQueue poll 再 offer 進 deque 的 last element;而有時需要從 maxQueue poll 出再 offer 進 deque 的 first element。
double median = 0;
Queue<Integer> minQueue;
Queue<Integer> maxQueue;
Deque<Integer> deque;
public MedianFinder() {
minQueue = new PriorityQueue<Integer>();
maxQueue = new PriorityQueue<Integer>(new Comparator<Integer>() {
public int compare(Integer n1, Integer n2) {
return n2 - n1;
}
});
deque = new ArrayDeque<Integer>();
}
public void addNum(int num) {
if (deque.size() == 0) {
deque.offerLast(num);
median = num;
} else {
if (num >= deque.peekLast()) {
minQueue.offer(num);
check();
} else if (num <= deque.peekFirst()) {
maxQueue.offer(num);
check();
} else {
median = num;
maxQueue.offer(deque.pollFirst());
minQueue.offer(deque.pollLast());
deque.offer(num);
}
}
}
public double findMedian() {
return median;
}
public void check() {
if (minQueue.size() != maxQueue.size()) {
if (minQueue.size() > maxQueue.size()) {
deque.offerLast(minQueue.poll());
} else {
deque.offerFirst(maxQueue.poll());
}
}
while (deque.size() > 2) {
maxQueue.offer(deque.pollFirst());
minQueue.offer(deque.pollLast());
}
if (deque.size() == 1) {
median = deque.peekLast();
} else {
median = (deque.peekLast() + deque.peekFirst()) / 2.0;
}
return;
}
不過參考了解法,發現不需要這麼麻煩,因為我們根本不需要維護中間那個 deque,畢竟那個 deque 是根據 maxQueue 或 minQueue 的 head 而來的。因此我們只需要確保這兩個 heap 的 size,要求 median 時再分別 peek 就可以
PriorityQueue<Integer> min = new PriorityQueue();
PriorityQueue<Integer> max = new PriorityQueue(1000, Collections.reverseOrder());
public void addNum(int num) {
max.offer(num);
min.offer(max.poll());
if (max.size() < min.size()){
max.offer(min.poll());
}
}
public double findMedian() {
if (max.size() == min.size()) return (max.peek() + min.peek()) / 2.0;
else return max.peek();
}