堆排序求中位数

二分查找求中位数需要数组有序，而且不能求动态数据的中位数，双堆就可以很好地解决这个问题。思路也比较简单，用大顶堆存较小的部分，用小顶堆存较大的部分，同时要维持两个堆的大小相差最多为1，这样求中位数的时候要么取某个堆的堆顶，要么求两个堆顶的平均值。

其实堆的问题最容易搞混的还是到底用什么堆存什么数据，比如求topk大的数就应该用小顶堆，求topk小的数就应该用大顶堆，还是要自己好好琢磨琢磨。

上代码：(规定maxHeap要么和minHeap一样大，要么比它多一个。看个人习惯，也可以minHeap多一个)

public class DualHeap { 
    private PriorityQueue<Integer> minHeap = new PriorityQueue<>(); // 不用 (o1, o2) -> o2 - o1，防整数相减溢出 // private PriorityQueue<Integer> maxHeap = new PriorityQueue<>((o1, o2) -> Integer.compare(o2, o1); private PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Comparator.reverseOrder()); public void addNum(int num) { 
    if (maxHeap.size() == minHeap.size()) { 
    minHeap.offer(num); maxHeap.offer(minHeap.poll()); } else { 
    maxHeap.offer(num); minHeap.offer(maxHeap.poll()); } } public double getMedian() { 
    if (maxHeap.size() == minHeap.size()) // 也是防整数相加溢出，leetcode还是挺爱出极端用例的 return maxHeap.peek() / 2.0 + minHeap.peek() / 2.0; else return maxHeap.peek(); } }

最开始双堆的大小都是0，即maxHeap.seize() == minHeap.size()。所以向minHeap中添加，然后再弹出堆顶，添加到maxHeap中，保证了maxHeap比比minHeap大1。

此时maxHeap.size() != minHeap.size()，应该让下一个元素添加到minHeap中：

此时maxHeap.size() == minHeap.seize()，下一个元素还应该向maxHeap中添加，还是和上面的步骤一样，想向哪个堆中添加，先去另一个堆中转一下，然后弹出最适合添加的数。所以后续的过程就不画了，最终两个堆应该是这样的：

所有数都添加完成后maxHeap.size() == minHeap.size()，所以返回中位数是二者堆顶的均值。如果size不相等，那肯定是maxHeap多一个，那就返回maxHeap.peek()就可以。

数据大小固定情况

上面的例子是可以一直添加数的，但如果是滑动窗口的题，比如今天的每日抑题leeetcode480，堆中数据的个数是固定的，那么在向堆中添加数(窗口右边界)的时候还要删一个数(窗口左边界)，所以上面的双堆代码就需要拓展，保证删除之后两个堆的大小还满足我们的要求。拓展之前先考虑一下从堆中删除数据的时候需要考虑哪些情况：

删之前，maxHeap.size() == minHeap.size()
- 要删的数在maxHeap中，删完之后maxHeap.sieze() < minHeap.size()，需要把minHeap的堆顶弹出，添加到maxHeap。
- 要删的数在minHeap中，删完之后maxHeap.size() = minHeap.size() + 1，不需要动。
删之前，maxHeap.size() = minHeap.size() + 1
- 要删的数在maxHeap中，删完之后maxHeap.size() == minHeap.size()，不需要动。
- 要删的数在minHeap中，删完之后maxHeap.size() = minHeap.size() + 2，需要把maxHeap的堆顶弹出，添加到minHeap中。

其实只有两种情况需要调整堆，那么代码就好写了：

public class DualHeap { 
    private PriorityQueue<Integer> minHeap = new PriorityQueue<>(); // 不用 (o1, o2) -> o2 - o1，防整数相减溢出 // private PriorityQueue<Integer> maxHeap = new PriorityQueue<>((o1, o2) -> Integer.compare(o2, o1); private PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Comparator.reverseOrder()); public void addNum(int num) { 
    if (maxHeap.size() == minHeap.size()) { 
    minHeap.offer(num); maxHeap.offer(minHeap.poll()); } else { 
    maxHeap.offer(num); minHeap.offer(maxHeap.poll()); } } public double getMedian() { 
    if (maxHeap.size() == minHeap.size()) // 也是防整数相加溢出，leetcode还是挺爱出极端用例的 return maxHeap.peek() / 2.0 + minHeap.peek() / 2.0; else return maxHeap.peek(); } public void removeNum(int num) { 
    if (num <= maxHeap.peek()) maxHeap.remove(num); else minHeap.remove(num); balance() } public void balance() { 
    if (maxHeap.size() > minHeap.size() + 1) minHeap.offer(maxHeap.poll()); else if (maxHeap.size() < minHeap.size()) maxHeap.offer(minHeap.poll()); } }

数据大小固定情况

知秋君

相关推荐