寻找中位数算法通用总结

寻找中位数我们通常可以从三个角度去思考，第一种方式利用双指针(前提是已经分别排序好)，第二种方式利用排序，第三种方式利用堆或者树。这里我们主要介绍利用堆或者树的解题模板。

P295 MedianFinder(难度中等)

解题思路，定义一个最大堆，定义一个最小堆，保证最大堆中的元素大于最小堆中的元素数量并且不超过一个。可以通过一段代码套用模板实现，先插入最大堆，然后将最大堆中的堆顶元素(最大值)插入到最小堆中；如果最大堆中的元素小于最小堆中的元素数量，将最小堆的堆顶元素(最小值)弹出到最大堆中：

public void addNum(int num) {
    maxHeap.offer(num);
    minHeap.offer(maxHeap.poll());
    if (maxHeap.size()<minHeap.size()) {
      maxHeap.offer(minHeap.poll());
    }
}

根据上面的方式构造的最大最小堆，我们只需要考虑堆顶的元素即可找出中位数，如果是偶数，那么中位数是两个堆顶元素的平均数，如果是奇数，那么中位数是最大堆的堆顶元素。这道题得到的代码如下：

package Design.P295;

import java.util.Collections;
import java.util.PriorityQueue;
import java.util.Queue;

/**
 * Example:
 *
 * Adding number 41
 * MaxHeap lo: [41]           // MaxHeap stores the largest value at the top (index 0)
 * MinHeap hi: []             // MinHeap stores the smallest value at the top (index 0)
 * Median is 41
 * =======================
 * Adding number 35
 * MaxHeap lo: [35]
 * MinHeap hi: [41]
 * Median is 38
 * =======================
 * Adding number 62
 * MaxHeap lo: [41, 35]
 * MinHeap hi: [62]
 * Median is 41
 * =======================
 * Adding number 4
 * MaxHeap lo: [35, 4]
 * MinHeap hi: [41, 62]
 * Median is 38
 * =======================
 * Adding number 97
 * MaxHeap lo: [41, 35, 4]
 * MinHeap hi: [62, 97]
 * Median is 41
 * =======================
 * Adding number 108
 * MaxHeap lo: [41, 35, 4]
 * MinHeap hi: [62, 97, 108]
 * Median is 51.5
 *
 *
 *
 *
 * addNum(2)
 * findMedian() -> 1.5
 * addNum(3)
 * findMedian() -> 2
 */
class MedianFinder {

    /** initialize your data structure here. */
    /** 定义一个最大堆和一个最小堆*/
    Queue<Integer> minHeap= new PriorityQueue<>();
    //传入一个倒序的比较器
    Queue<Integer> maxHeap= new PriorityQueue<>(Collections.reverseOrder());
    public MedianFinder() {

    }
    
    public void addNum(int num) {
        maxHeap.offer(num);
        minHeap.offer(maxHeap.poll());
        if (maxHeap.size()<minHeap.size()) {
            maxHeap.offer(minHeap.poll());
        }
    }
    
    public double findMedian() {
        int minSize=minHeap.size();
        int maxSize=maxHeap.size();

        if (((maxSize+minSize)&1)==0) {
            return (minHeap.peek()+maxHeap.peek())/2.0;
        }else {
            return maxHeap.peek()/1.0;
        }
    }

    public static void main(String[] args) {
        MedianFinder medianFinder=new MedianFinder();
        medianFinder.addNum(1);
        medianFinder.addNum(2);
        System.out.println(medianFinder.findMedian()); //-> 1.5
        medianFinder.addNum(3);
        System.out.println(medianFinder.findMedian()); //-> 2
        medianFinder.addNum(5);
        System.out.println(medianFinder.findMedian());
        medianFinder.addNum(4);
        System.out.println(medianFinder.findMedian());
    }
}

/**
 * Your MedianFinder object will be instantiated and called as such:
 * MedianFinder obj = new MedianFinder();
 * obj.addNum(num);
 * double param_2 = obj.findMedian();
 */

P480 MedianSlidingWindow(难度困难)

这道题不同的是如果只用最大堆和最小堆来实现的话，在删除节点的时候复杂度过高了，为了避免删除节点造成的时间复杂度提升，我们可以改用TreeMap二叉搜索树来实现，将删除节点的复杂度从n减少到了logn。我们套用构建最大树和最小树的模板，发现是极其相似的：

for (int i = 0; i < k; i++) {
    left.add(i);
    right.add(left.pollLast());
    if (left.size()<right.size()) {
      left.add(right.pollFirst());
    }
}

其实这边采用两个和采用一个TreeMap是一样的，只不过始终需要保持左子树和右子树的节点数量差左边大于右边不超过一个。参考一个更加通俗的本题的解题模板。

package SlidingWindow.P480;

import java.util.Comparator;
import java.util.TreeSet;

/**
 * 这个版本更加通俗易懂
 * @autor yeqiaozhu.
 * @date 2020-01-02
 */
public class UsingTreeSetClear {
    public double[] medianSlidingWindow(int[] nums, int k) {
        double[] result = new double[nums.length - k + 1];
        TreeSet<Integer> left = getSet(nums);
        TreeSet<Integer> right = getSet(nums);
        for (int i = 0; i < nums.length; i++) {
            if (left.size() <= right.size()) {
                right.add(i);
                int m = right.first();
                right.remove(m);
                left.add(m);
            } else {
                left.add(i);
                int m = left.last();
                left.remove(m);
                right.add(m);
            }


            if (left.size() + right.size() == k) {
                double med;
                if (left.size() == right.size())
                    med = ((double) nums[left.last()] + nums[right.first()]) / 2;
                else if (left.size() < right.size())
                    med = nums[right.first()];
                else
                    med = nums[left.last()];

                int start = i - k + 1;
                result[start] = med;

                if (!left.remove(start))
                    right.remove(start);
            }
        }
        return result;
    }
    private static TreeSet<Integer> getSet(int[] nums) {
        return new TreeSet<>(new Comparator<Integer>(){
            public int compare(Integer a, Integer b) {
                return nums[a] == nums[b] ? a - b : nums[a] < nums[b] ? -1 : 1;
            }
        });
    }
}

总结

构造最大最小堆/树可以参考其中一个模板。模块的开头已经列出。