[堆 二叉搜索树] 295. 数据流的中位数（排序法 → 二叉搜索树法（手写BST+查找第k小元素）、大小顶堆法（设计出入堆规则））

295. 数据流的中位数

题目链接：https://leetcode-cn.com/problems/find-median-from-data-stream/

---

分类：

• 数学：中位数 = 有序列表中间的数或中间两个数的平均值，根据中位数特点设计addNum和findMedian函数；
• 排序：每输入一个数据就对整个列表排序（Collections.sort），再用下标取中位数；（思路1）
• 二叉搜索树：数据流存入二叉搜索树，再查找第n/2+1小的数字（一大知识点）；（思路2，用来练习手写二叉搜索树）
• 堆：设置大小顶堆，设计出入堆规则，保证两个堆的顶部保存中间两个数；（思路3，首推）

---

题目分析

中位数是有序列表中间的数，如果数据列表有序，则要找中位数很容易，前两个思路是对整个列表做排序，再寻找中位数；思路3则只关注中位数本身，其他元素只保持大致有序。

注意：在分析时间复杂度方面，官方题解是分别对调用一次addNum函数和调用一次findMedian函数分析时间复杂度，而不是分析整个寻找中位数的过程。

进阶（待补充）:

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

   元素取值范围固定，是否可以考虑计数排序？

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

   是否可以考虑使用map？

思路1：排序

使用一个列表存放输入的数据，每加入一个数字，就对列表进行排序，根据下标找出中位数。

实现比较简单，不再赘述。

• 时间复杂度：注意这里对时间复杂度的分析是把addNum和findMedian作为两个单独的操作分析的。
  • addNum：每调用一次addNum加入一个新数字就需要排序一次，所以需要O(NlogN)；
  • findMedian：O(1)。

思路2：二叉搜索树法（练习手写BST，比较麻烦）

思路2实际上和思路1一样也是排序，只是把数据存入二叉搜索树中再做查找，这就需要我们手动实现二叉搜索树，需要编写的有：节点类TreeNode，添加元素函数addNum()和查找第k小的函数searchKth()。

• addNum：新数据在二叉搜索树中做二分查找寻找合适的位置后插入；
• findMedian：
  • 如果数据个数为奇数，则在二叉搜索树中查找第n/2个元素（下标从0开始计）；
  • 如果数据个数为偶数，则在二叉搜索树中查找第n/2和第n/2-1个元素。

关键问题：如何在二叉搜索树上查找第 k 小的元素（k从1开始计）？

---

方法1：中序遍历 + 全局变量

对BST做中序遍历能够得到有序序列，第k小的元素就是中序遍历到的第k个节点，所以设置一个全局变量count记录中序遍历访问过的节点个数，当count==k时，说明找到第k小的元素。

searchKth 实现代码：

    //version1:dfs+使用全局变量记录遍历个数，优点：代码简单；缺点：只能使用一次
    int cnt = 0;
    int res = 0;
    public void searchKth(int k, TreeNode root){
        if(root == null) return;

        searchKth(k, root.left);
        cnt++;
        if(cnt == k) {
            res = root.val;
            return;
        }
        searchKth(k, root.right);
    }

• 优点：代码实现容易；
• 缺点：只能使用一次，调用一次查找函数就会修改到全局变量，导致下一次调用时全局变量就不再适用了。

方法2（采用）：拿左子树节点个数和k比较
对于root，计算它的左子树节点个数leftCount：

• 如果leftCount + 1 == k，说明当前root就是第k小的元素；
• 如果leftCount + 1 > k，说明第i小的元素在root的左子树，进入左子树继续递归查找；
• 如果leftCount + 1 < k，说明第i小的元素在root的右子树，需要进入右子树查找，但在右子树中我们要查找的就不再是第k小的元素，而是排除了左子树和根节点后，第k-leftCount-1小的元素。

其中，如何计算一棵树的节点个数也是一个问题：

version1：对树做一次遍历，统计节点个数，缺点是每个节点都需要做同样的计算，很耗时；

searchKth 实现代码：

	//查找第k小的元素(k从1开始计)：通过左子树节点个数来判断
    public int searchKth(int k, TreeNode root){
        int leftCount = getNum(root.left);
        if(leftCount + 1 == k){
            return root.val;
        }
        else if(leftCount + 1 > k){
            return searchKth(k, root.left);
        }
        else{
            return searchKth(k - leftCount - 1, root.right);
        }
    }
    //获取当前节点的所有子孙节点个数（包括节点本身）
    public int getNum(TreeNode root){
        if(root == null) return 0;
        else{
            return getNum(root.left) + getNum(root.right) + 1;
        }
    }

• 在本题中version1会超时。

version2（采用）：在树节点中增加一个size成员变量，记录以当前节点为根的树有多少个节点（包括当前节点本身），每个新创建的节点的size=1，每向BST添加一个新节点，在寻找合适位置过程中，经过的每一个节点的size都要+1。

searchKth 实现代码：（需要修改到addNum，TreeNode，search，这里只列了search）

        
	//查找第k小的元素(k从1开始计)：通过左子树节点个数来判断
    public int searchKth(int k, TreeNode root){
        int leftCount = root.left == null ? 0 : root.left.size;//记得对左子树判空
        if(leftCount + 1 == k){
            return root.val;
        }
        else if(leftCount + 1 > k){
            return searchKth(k, root.left);
        }
        else{
            return searchKth(k - leftCount - 1, root.right);
        }
    }

思路2整体实现代码（采用方法2-version2）

class MedianFinder {
    //树节点
    class TreeNode{
        int val;
        TreeNode left;
        TreeNode right;
        int size;//记录以该节点为根的所有子孙节点个数（包括节点本身）
        public TreeNode(int val){
            this.val = val;
            this.size = 1;//创建一个新节点，size=1
        }
    }
    
    TreeNode root;//根节点
    int count;//统计总个数

    /** initialize your data structure here. */
    public MedianFinder() {
        root = null;
        count = 0;
    }
    
    public void addNum(int num) {
        //如果root为null，则插入根节点处
        if(root == null){
            root = new TreeNode(num);
        }
        else{
            addNum(num, root);
        }
        count++;
    }
    //添加元素的辅助函数
    private TreeNode addNum(int num, TreeNode node){
        if(node == null) return new TreeNode(num);
        if(num <= node.val){
            node.left = addNum(num, node.left);
        }
        else{
            node.right = addNum(num, node.right);
        }
        node.size++;//寻找合适位置过程中经过的点的size都要加1
        return node;
    }
    
    public double findMedian() {
        if(count % 2 == 1){
            //中位数的下标是count/2，但按大小排是第count/2+1个元素（大小从1开始计）
            return searchKth(count / 2 + 1, root);
        }
        else return (searchKth(count / 2, root) + searchKth(count / 2 + 1, root)) / 2.0;
    }
    //查找第index小的元素(index从1开始计)：通过左子树节点个数来判断
    public int searchKth(int k, TreeNode root){
        int leftCount = root.left == null ? 0 : root.left.size;//记得对左子树判空
        if(leftCount + 1 == k){
            return root.val;
        }
        else if(leftCount + 1 > k){
            return searchKth(k, root.left);
        }
        else{
            return searchKth(k - leftCount - 1, root.right);
        }
    }
}

• 时间复杂度：思路2的addNum和findMedian的时间复杂度近似，都是基于二叉搜索树的查找：平均需要O(logN)，最差情况下需要O(N)。

思路3：大小堆法（推荐）

思路1的时间消耗在对整个列表的排序上，实际上我们只需要关心整个列表中的中位数就好，其他数字并不需要严格排序，由此就想到堆这一数据结构。

我们先观察一组有序数据的中位数特点：

例如：nums=[1,2,3,4,5,6]

nums元素个数是偶数，所以中位数是(3+4)/2，所以如果要寻找中位数只需要关注3和4，至于3左边的1,2和4右边的5,6的排列是否有序对中位数并没有影响，我们只需要保证3,4在正确的位置即可。

因为3是[1,2,3]中的最大值，4是[4,5,6]中的最小值，所以我们可以创建一个大顶堆存放[1,2,3]，一个小顶堆存放[4,5,6]，要计算中位数，先判断两个堆的大小关系：

• 如果两个堆大小相等，说明数据个数为偶数，中位数=两个堆的顶部(3+4)/2；
• 如果两个堆大小不等，说明数据个数为奇数，中位数=较大堆的顶部。

简单来说，两个堆的顶部保存的是“有序”数据流n的第n/2和第n/2-1个元素（下标从0开始计），但输入的数据流是无序的，如何设计堆的出入规则，确保两个堆满足“顶部保存的是有序数据流n的第n/2和第n/2-1个元素”？

设计出入堆规则需要满足两个条件：

1. 确保小顶堆的顶部 > 大顶堆的顶部；
2. 两个堆的大小差值不能超过1。为了使算法更清晰，我们规定小顶堆的大小 >= 大顶堆的大小，这样如果数字个数是偶数，则中位数=(小顶堆顶部+大顶部顶部)/2，如果数字个数是奇数，则中位数=小顶堆顶部。

出入堆规则：（addNum）
对于一个新数据：

• 如果两个堆都为空，则加入大顶堆；
• 如果两个堆不全为空，则大顶堆此时一定不为空，则拿大顶堆顶部和新数据比较：
  • 如果新数据 <= 大顶堆顶部，则加入大顶堆；
  • 如果新数据 > 大顶堆顶部，则加入小顶堆。

加入之后需要做一次堆的大小平衡：（关键点）

• 如果此时大顶堆size - 小顶堆size > 1，则弹出大顶堆顶部给小顶堆，将size差值降低为1；
• 如果此时小顶堆size - 大顶堆size >=1，则弹出小顶堆顶部给大顶堆，确保大顶堆的size总是>=小顶堆的size。
• 其他情况保持不变。

获取中位数：（findMedian）
判断入堆个数的奇偶性：

• 如果为偶数，则中位数=(小顶堆顶部+大顶部顶部)/2；
• 如果为奇数，则中位数=大顶堆顶部。

实现代码：

class MedianFinder {
    PriorityQueue<Integer> maxHeap;
    PriorityQueue<Integer> minHeap;
    int count;
    /** initialize your data structure here. */
    public MedianFinder() {
        maxHeap = new PriorityQueue<>((x, y) -> y - x);//lambda初始化大顶堆
        minHeap = new PriorityQueue<>();//默认是小顶堆
        count = 0;
    }
    
    public void addNum(int num) {
        //如果两个堆都为空，则加入小顶堆
        if(count == 0) maxHeap.offer(num);
        //如果两个堆不全为空，则小顶堆一定不为空（出入规则造成的）
        else{
            if(num <= maxHeap.peek()) maxHeap.offer(num);
            else minHeap.offer(num);
        }
        count++;
        //堆的大小平衡处理：
        //如果小顶堆size-大顶堆size>1，则弹出小顶堆顶部给大顶堆，将size差值降低为1；
        if(maxHeap.size() - minHeap.size() > 1) minHeap.offer(maxHeap.poll());
        //如果此时大顶堆size-小顶堆size>=1，则弹出大顶堆顶部给小顶堆，确保小顶堆的size总是>=大顶堆的size
        else if(minHeap.size() - maxHeap.size() >= 1) maxHeap.offer(minHeap.poll());
        //其他情况保持不变。
    }
    
    public double findMedian() {
        if(count % 2 == 0) return (minHeap.peek() + maxHeap.peek()) / 2.0;
        else return maxHeap.peek();
    }
}

• 时间复杂度：

  • addNum：每一个新元素都会加入到大顶堆或小顶堆，加入堆中就会引起堆化处理，需要O(logN)；同时还可能引起两个堆大小平衡操作，但根据算法流程的设计，大小平衡处理每次最多只会移动一个元素，所以需要O(1)，整体时间复杂度为O(logN)。
  • findMedian：O(1).

• 空间复杂度：两个堆一起存放数据流的所有数据，所以占用空间O(N)。