【暴力, 二分查找, 优先队列】数据流的中位数

题目

295. 数据流的中位数

---

中位数是有序列表中间的数。如果列表长度是偶数，中位数则是中间两个数的平均值。

例如，

[2,3,4] 的中位数是 3

[2,3] 的中位数是 (2 + 3) / 2 = 2.5

设计一个支持以下两种操作的数据结构：

• void addNum(int num) - 从数据流中添加一个整数到数据结构中。
• double findMedian() - 返回目前所有元素的中位数。

示例：

addNum(1)
addNum(2)
findMedian() -> 1.5
addNum(3) 
findMedian() -> 2

进阶:

• 如果数据流中所有整数都在 0 到 100 范围内，你将如何优化你的算法？
• 如果数据流中 99% 的整数都在 0 到 100 范围内，你将如何优化你的算法？

解题

方法一：暴力 排序

思路

按照题意模拟，使用一个集合(numList)存储数字，添加数字的时候直接加入集合，获取中位数的时候把集合排序再找出中位数返回。

代码

class MedianFinder {
    private List<Integer> numList;

    public MedianFinder() {
        numList = new ArrayList<>();
    }
    
    public void addNum(int num) {
        numList.add(num);
    }
    
    public double findMedian() {
        numList.sort((a, b) -> a - b);
        int len = numList.size();
        if ((len & 1) == 0) {
            return (numList.get(len / 2) + numList.get(len / 2 - 1)) / 2.0;
        } else return numList.get(len / 2);
    }
}

方法二：二分查找

思路

使用一个集合(numList)存储数字，添加数字的时候使用二分查找到 numList 中小于这个数字中最大的那个数字的索引，把新添加的数字插入在那个索引后面。
获取中位数的时候集合就已经是排好序的了，直接找中位数返回。

代码

class MedianFinder {
    private List<Integer> numList;

    public MedianFinder() {
        numList = new ArrayList<>();
    }
    
    public void addNum(int num) {
        int start = 0, end = numList.size() - 1;
        while (start <= end) {
            int mid = start + ((end - start) >> 1);
            if (numList.get(mid) <= num) start = mid + 1;
            else end = mid - 1;
        }
        numList.add(end + 1, num);
    }
    
    public double findMedian() {
        int len = numList.size();
        if ((len & 1) == 0) {
            return (numList.get(len / 2) + numList.get(len / 2 - 1)) / 2.0;
        } else return numList.get(len / 2);
    }
}

优化

对部分测试用例进行判断优化

class MedianFinder {
    private List<Integer> numList;

    public MedianFinder() {
        numList = new ArrayList<>();
    }
    
    public void addNum(int num) {
        int start = 0, end = numList.size() - 1;
        if (end == -1) numList.add(num);
        else if (num <= numList.get(0)) numList.add(0, num);
        else if (num >= numList.get(end)) numList.add(num);
        else {
            while (start <= end) {
                int mid = start + ((end - start) >> 1);
                if (numList.get(mid) < num) start = mid + 1;
                else end = mid - 1;
            }
            numList.add(end + 1, num);
        }
    }
    
    public double findMedian() {
        int len = numList.size();
        if ((len & 1) == 0) {
            return (numList.get(len / 2) + numList.get(len / 2 - 1)) / 2.0;
        } else return numList.get(len / 2);
    }
}

方法三：优先队列

思路

维护两个优先队列来维护数据流数据：一个大根堆(最大优先队列)(leftPriQueue)用来存放列表左半边的数字，一个小根堆(最小优先队列)(rightPriQueue)用来存放列表右半边的数字，其中它们有如下关系：

• 当数据流元素数量为偶数：leftPriQueue.size() == rightPriQueue.size() ，此时动态中位数为两者堆顶元素的平均值；
• 当数据流元素数量为奇数：leftPriQueue.size() == rightPriQueue.size() + 1，此时动态中位数为 leftPriQueue 的堆顶原数。

在把一个数 num 添加进数据流中时，要分情况讨论：

• $num \leq \max\{leftPriQueue\}$ 时：
  num 小于等于中位数，应该将该数添加到 leftPriQueue 中。新的中位数将小于等于原中位数，因此可能需要将 leftPriQueue 中的原中位数也就是最大数移动到 rightPriQueue。
• $num > \max\{leftPriQueue\}$ 时：
  num 大于中位数，应该将该数添加到 rightPriQueue 中。新的中位数将大于等于原中位数，因此可能需要将 rightPriQueue 中的原中位数也就是最小数移动到 leftPriQueue。
• 特别地，当数据流中还没有数时，应该把 num 添加到 leftPriQueue 中。

代码

class MedianFinder {
    private Queue<Integer> leftPriQueue;
    private Queue<Integer> rightPriQueue;

    public MedianFinder() {
        leftPriQueue = new PriorityQueue<>((a, b) -> b - a);
        rightPriQueue = new PriorityQueue<>((a, b) -> a - b);
    }
    
    public void addNum(int num) {
        if (leftPriQueue.isEmpty() || num <= leftPriQueue.peek()) {
            leftPriQueue.offer(num);
            if (leftPriQueue.size() - rightPriQueue.size() > 1) {
                rightPriQueue.offer(leftPriQueue.poll());
            }
        } else {
            rightPriQueue.offer(num);
            if (rightPriQueue.size() > leftPriQueue.size()) {
                leftPriQueue.offer(rightPriQueue.poll());
            }
        }
    }
    
    public double findMedian() {
        return leftPriQueue.size() == rightPriQueue.size() ?
            (leftPriQueue.peek() + rightPriQueue.peek()) / 2.0 :
            leftPriQueue.peek();
    }
}