双堆求中位数(java)

二分查找求中位数需要数组有序，而且不能求动态数据的中位数，双堆就可以很好地解决这个问题。思路也比较简单，用大顶堆存较小的部分，用小顶堆存较大的部分，同时要维持两个堆的大小相差最多为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());
    }
}