P159 在数据流中查找中位数

描述#

这个题目说的是，你要设计一个数据结构用于查找数据流中的中位数。它要支持以下两个操作，addNum 用于向数据结构中添加一个整数。findMedian 用于查找该数据结构中目前所有整数的中位数。

中位数的定义是将序列排序后正中间的数字，如果序列的长度是偶数，则取正中间两个数字的平均数。

代码#

public class AlgoCasts {

  public class MedianFinderArray {

    private final List<Integer> data = new ArrayList<>();

    // Time: O(n)
    public void addNum(int num) {
      int idx = data.size() - 1;
      while (idx >= 0 && data.get(idx) > num) --idx;
      data.add(idx+1, num);
    }

    // Time: O(1)
    public double findMedian() {
      int n = data.size();
      if ((n & 1) == 1) {
        return data.get(n/2);
      } else {
        return (data.get(n/2-1) + data.get(n/2)) / 2.0;
      }
    }
  }

  public class MedianFinderHeap {

    private final Queue<Integer> minHeap = new PriorityQueue<>();
    private final Queue<Integer> maxHeap = new PriorityQueue<>(Collections.reverseOrder());

    // Time: O(log(n))
    public void addNum(int num) {
      minHeap.add(num);
      maxHeap.add(minHeap.poll());
      if (maxHeap.size() - minHeap.size() > 1)
        minHeap.add(maxHeap.poll());
    }

    // Time: O(1)
    public double findMedian() {
      if (maxHeap.size() > minHeap.size()) {
        return maxHeap.peek();
      } else {
        return (maxHeap.peek() + minHeap.peek()) / 2.0;
      }
    }
  }

}