数据流的中位数

题目要求

设计一个 MedianFinder 类，该类能够：

1. 通过 addNum(int num) 方法接收一个数据流中的整数 num 并将其加入到内部数据结构中。
2. 通过 findMedian() 方法计算并返回当前所有加入的整数的中位数。

中位数的定义是：在一个有序的整数列表中，位于中间位置的数。如果列表的长度是奇数，则中位数是中间的数；如果列表的长度是偶数，则中位数是中间两个数的平均值。

解题思路

为了实现 MedianFinder 类，我们需要一个有效的内部数据结构来存储输入的整数，并且能够快速计算中位数。以下是解决这个问题的思路：

1. 数据存储：使用两个优先队列（通常实现为堆），一个最大堆来存储当前所有数中较小的一半，一个最小堆来存储较大的一半。这样可以保证最大堆的最大元素和最小堆的最小元素是所有数中间的两个数。

2. 平衡堆：为了快速找到中位数，需要保持两个堆的大小平衡（或者最大堆比最小堆多一个元素，这样当元素总数为奇数时，最大堆的顶部就是中位数）。

3. 添加元素：当添加一个新元素时，可以先将其加入到最大堆，然后将最大堆的顶部元素移动到最小堆。这样做是为了保持元素分布的平衡。如果最大堆的大小超过了最小堆的大小超过 1，就将最小堆的顶部元素移动到最大堆。

4. 计算中位数：如果两个堆的大小相等，则中位数是两个堆顶元素的平均值；如果最大堆的元素多，则中位数是最大堆的堆顶元素。

通过这种方式，addNum 操作可以在对数时间内完成，findMedian 操作可以在常数时间内完成，因为只涉及到获取堆顶元素和计算平均值。

Golang 版本

import (
    "container/heap"
)

// An IntHeap is a min-heap of ints.
type IntHeap []int

func (h IntHeap) Len() int           { return len(h) }
func (h IntHeap) Less(i, j int) bool { return h[i] < h[j] }
func (h IntHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *IntHeap) Push(x interface{}) {
    *h = append(*h, x.(int))
}

func (h *IntHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[0 : n-1]
    return x
}

// An IntMaxHeap is a max-heap of ints.
type IntMaxHeap struct {
    IntHeap
}

func (h IntMaxHeap) Less(i, j int) bool { return h.IntHeap[i] > h.IntHeap[j] }

// MedianFinder holds two heaps.
type MedianFinder struct {
    minHeap IntHeap
    maxHeap IntMaxHeap
}

/** initialize your data structure here. */
func Constructor() MedianFinder {
    return MedianFinder{
        minHeap: make(IntHeap, 0),
        maxHeap: IntMaxHeap{make(IntHeap, 0)},
    }
}

func (this *MedianFinder) AddNum(num int) {
    heap.Push(&this.maxHeap, num)
    heap.Push(&this.minHeap, heap.Pop(&this.maxHeap))
    if this.minHeap.Len() > this.maxHeap.Len() {
        heap.Push(&this.maxHeap, heap.Pop(&this.minHeap))
    }
}

func (this *MedianFinder) FindMedian() float64 {
    if this.maxHeap.Len() > this.minHeap.Len() {
        return float64(this.maxHeap.IntHeap[0])
    }
    return float64(this.minHeap.IntHeap[0]+this.maxHeap.IntHeap[0]) / 2.0
}

/**
 * Your MedianFinder object will be instantiated and called as such:
 * obj := Constructor();
 * obj.AddNum(num);
 * param_2 := obj.FindMedian();
 */

在这段代码中，我们定义了两个类型 IntHeap 和 IntMaxHeap 来实现最小堆和最大堆。MedianFinder 结构体包含了一个最小堆和一个最大堆，用于存储数据流中的数字。AddNum 方法将一个新的数字加入到合适的堆中，并保持两个堆的平衡。FindMedian 方法根据两个堆的顶部元素计算出中位数。

C++版本

#include <queue>
#include <vector>
#include <functional>

class MedianFinder {
private:
    // 大根堆，存储较小的一半元素
    std::priority_queue<int, std::vector<int>, std::less<int>> maxHeap;
    // 小根堆，存储较大的一半元素
    std::priority_queue<int, std::vector<int>, std::greater<int>> minHeap;

public:
    /** initialize your data structure here. */
    MedianFinder() {

    }

    void addNum(int num) {
        // 先加入大根堆，再将大根堆的最大元素加入小根堆
        maxHeap.push(num);
        minHeap.push(maxHeap.top());
        maxHeap.pop();

        // 如果小根堆元素多，移动到大根堆
        if (minHeap.size() > maxHeap.size()) {
            maxHeap.push(minHeap.top());
            minHeap.pop();
        }
    }

    double findMedian() {
        // 如果两个堆的大小相等，则取平均值；否则取大根堆的顶部元素
        if (maxHeap.size() > minHeap.size()) {
            return maxHeap.top();
        } else {
            return (maxHeap.top() + minHeap.top()) * 0.5;
        }
    }
};

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

在这段 C++ 代码中，我们使用了 STL 中的 priority_queue 来实现最大堆和最小堆。MedianFinder 类中包含两个 priority_queue，一个作为最大堆（maxHeap），用于存储较小的一半元素，另一个作为最小堆（minHeap），用于存储较大的一半元素。addNum 方法负责将新元素加入到适当的堆中，并保持两个堆的平衡。findMedian 方法则根据两个堆的顶部元素来计算当前的中位数。

Python 版本

from heapq import *

class MedianFinder:

    def __init__(self):
        # 小顶堆，存储较大的一半元素
        self.minHeap = []
        # 大顶堆，存储较小的一半元素，我们通过存储相反数来实现
        self.maxHeap = []

    def addNum(self, num):
        # 先加入大顶堆，再将大顶堆的最大元素加入小顶堆
        heappush(self.maxHeap, -num)
        heappush(self.minHeap, -heappop(self.maxHeap))

        # 如果小顶堆元素多，移动到大顶堆
        if len(self.minHeap) > len(self.maxHeap):
            heappush(self.maxHeap, -heappop(self.minHeap))

    def findMedian(self):
        # 如果两个堆的大小相等，则取平均值；否则取大顶堆的顶部元素
        if len(self.minHeap) == len(self.maxHeap):
            return (-self.maxHeap[0] + self.minHeap[0]) / 2.0
        else:
            return -self.maxHeap[0]

# Your MedianFinder object will be instantiated and called as such:
# obj = MedianFinder()
# obj.addNum(num)
# param_2 = obj.findMedian()

在这段 Python 代码中，我们使用了 Python 内置的heapq模块来实现最小堆。为了实现最大堆，我们将元素的符号取反后存入最小堆中。MedianFinder类中包含两个列表，一个表示最小堆（minHeap），一个表示最大堆（maxHeap）。addNum方法用于将新元素添加到合适的堆中，并保持两个堆的平衡。findMedian方法根据两个堆的顶部元素来计算当前的中位数。

Java 版本

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

public class MedianFinder {

    private PriorityQueue<Integer> minHeap;
    private PriorityQueue<Integer> maxHeap;

    /** initialize your data structure here. */
    public MedianFinder() {
        // 小顶堆，存储较大的一半元素
        minHeap = new PriorityQueue<>();
        // 大顶堆，存储较小的一半元素
        maxHeap = new PriorityQueue<>(Collections.reverseOrder());
    }

    public void addNum(int num) {
        // 先加入大顶堆，再将大顶堆的最大元素加入小顶堆
        maxHeap.offer(num);
        minHeap.offer(maxHeap.poll());

        // 如果小顶堆元素多，移动到大顶堆
        if (minHeap.size() > maxHeap.size()) {
            maxHeap.offer(minHeap.poll());
        }
    }

    public double findMedian() {
        // 如果两个堆的大小相等，则取平均值；否则取大顶堆的顶部元素
        if (minHeap.size() == maxHeap.size()) {
            return (minHeap.peek() + maxHeap.peek()) / 2.0;
        } else {
            return maxHeap.peek();
        }
    }
}

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

在这段 Java 代码中，我们使用了 Java 的PriorityQueue类来实现最小堆和最大堆。PriorityQueue默认实现的是最小堆，为了实现最大堆，我们在构造时传入了Collections.reverseOrder()比较器。MedianFinder类中包含两个PriorityQueue，一个表示最小堆（minHeap），一个表示最大堆（maxHeap）。addNum方法用于将新元素添加到合适的堆中，并保持两个堆的平衡。findMedian方法根据两个堆的顶部元素来计算当前的中位数。

Rust 版本

use std::collections::BinaryHeap;
use std::cmp::Reverse;

struct MedianFinder {
    min_heap: BinaryHeap<Reverse<i32>>, // 小顶堆，存储较大的一半元素
    max_heap: BinaryHeap<i32>,          // 大顶堆，存储较小的一半元素
}

impl MedianFinder {
    /** initialize your data structure here. */
    fn new() -> Self {
        MedianFinder {
            min_heap: BinaryHeap::new(),
            max_heap: BinaryHeap::new(),
        }
    }

    fn add_num(&mut self, num: i32) {
        // 先加入大顶堆，再将大顶堆的最大元素加入小顶堆
        self.max_heap.push(num);
        let max_heap_top = self.max_heap.pop().unwrap();
        self.min_heap.push(Reverse(max_heap_top));

        // 如果小顶堆元素多，移动到大顶堆
        if self.min_heap.len() > self.max_heap.len() {
            let Reverse(min_heap_top) = self.min_heap.pop().unwrap();
            self.max_heap.push(min_heap_top);
        }
    }

    fn find_median(&self) -> f64 {
        if self.min_heap.len() == self.max_heap.len() {
            // 如果两个堆的大小相等，则取平均值
            let Reverse(min_heap_top) = *self.min_heap.peek().unwrap();
            let max_heap_top = *self.max_heap.peek().unwrap();
            (min_heap_top as f64 + max_heap_top as f64) / 2.0
        } else {
            // 否则取大顶堆的顶部元素
            *self.max_heap.peek().unwrap() as f64
        }
    }
}

fn main() {
    let mut finder = MedianFinder::new();
    finder.add_num(1);
    finder.add_num(2);
    println!("Median is {}", finder.find_median()); // 输出中位数 1.5
    finder.add_num(3);
    println!("Median is {}", finder.find_median()); // 输出中位数 2
}

在这段 Rust 代码中，我们使用了 Rust 标准库中的BinaryHeap来实现最大堆。为了实现最小堆，我们使用了Reverse包装器来改变元素的排序方式。MedianFinder结构体包含两个BinaryHeap，一个表示最小堆（min_heap），一个表示最大堆（max_heap）。add_num方法用于将新元素添加到合适的堆中，并保持两个堆的平衡。find_median方法根据两个堆的顶部元素来计算当前的中位数。在main函数中，我们创建了一个MedianFinder实例，并添加了一些数字来演示如何计算中位数。

总结

上述 Rust 版本的解法中，我们实现了一个MedianFinder结构体，它使用两个二叉堆来有效地计算动态整数集合的中位数。这种方法的核心思想是保持两个堆的元素数量平衡或接近平衡，这样中位数就可以容易地从堆顶元素中得到。

• 最大堆（max_heap）：存储当前所有元素中较小的一半，使用 Rust 的BinaryHeap直接实现。
• 最小堆（min_heap）：存储当前所有元素中较大的一半，通过对元素应用Reverse包装器来逆转BinaryHeap的默认行为，从而实现最小堆。

添加元素（add_num方法）：

1. 新元素首先被添加到最大堆中。
2. 然后将最大堆的顶部元素（即最大元素）移动到最小堆中。
3. 如果最小堆的元素数量超过了最大堆，将最小堆的顶部元素（即最小元素）移回最大堆。 这样做可以保证最大堆的所有元素都小于或等于最小堆的所有元素，并且两个堆的大小差不会超过 1。

计算中位数（find_median方法）：

• 如果两个堆的大小相等，中位数是两个堆顶元素的平均值。
• 如果两个堆的大小不等，中位数是元素较多的堆的顶部元素。

这种方法的时间复杂度是：

• 添加元素：O(log n)，因为需要将元素插入堆中。
• 查找中位数：O(1)，因为直接访问堆顶元素。

这个解法适用于数据流中位数的场景，因为它可以在接收每个新元素后立即计算中位数。