295. 数据流的中位数

题目描述

把一组整数排好序后，位于正中间的数就是中位数。数的个数为偶数时，取中间两个数的平均值。例如 [2,3,4] 的中位数是 3，[2,3] 的中位数是 (2 + 3) / 2 = 2.5。

设计一个 MedianFinder 类，接收一串不断到来的整数，并能随时给出中位数。

MedianFinder()：初始化对象。void addNum(int num)：加入一个整数 num。double findMedian()：返回目前已加入的所有整数的中位数，与真实值误差在 10^-5 以内即可。

输入由方法名数组和参数数组组成，输出按顺序给出每次调用的返回值，无返回值的调用记为 null。

示例 1

输入：操作 = ["MedianFinder","addNum","addNum","findMedian","addNum","findMedian"]
参数 = [[],[1],[2],[],[3],[]]
输出：[null,null,null,1.5,null,2.0]

解释：加入 1、2 后中位数是 (1 + 2) / 2 = 1.5；再加入 3 后，[1,2,3] 的中位数是 2.0。

数据范围

-10^5 <= num <= 10^5
每次调用 findMedian 时，结构里至少已有一个数。
addNum 和 findMedian 合计最多调用 5 * 10^4 次。

解题思路：对顶双堆

用大顶堆保存较小的一半，用小顶堆保存较大的一半；插入后平衡两个堆。

为什么正确

大顶堆的所有元素不大于小顶堆的所有元素，且两堆大小相差不超过 1。

复杂度分析

addNum O(log n)，findMedian O(1)，空间 O(n)。

易错点

Python heapq 只有小顶堆，左堆要存负数；取中位数时记得还原符号。

参考代码

Python、Java、C++、Go、JavaScript、TypeScript 六种语言，逻辑一致。

import heapq

class MedianFinder:
    def __init__(self):
        self.left = []
        self.right = []
    def addNum(self, num: int) -> None:
        heapq.heappush(self.left, -num)
        heapq.heappush(self.right, -heapq.heappop(self.left))
        if len(self.right) > len(self.left):
            heapq.heappush(self.left, -heapq.heappop(self.right))
    def findMedian(self) -> float:
        if len(self.left) > len(self.right):
            return float(-self.left[0])
        return (-self.left[0] + self.right[0]) / 2

class MedianFinder {
    private PriorityQueue<Integer> left = new PriorityQueue<>(Collections.reverseOrder());
    private PriorityQueue<Integer> right = new PriorityQueue<>();

    public MedianFinder() {}

    public void addNum(int num) {
        left.offer(num);
        right.offer(left.poll());
        if (right.size() > left.size())
            left.offer(right.poll());
    }

    public double findMedian() {
        if (left.size() > right.size())
            return left.peek();
        return (left.peek() + right.peek()) / 2.0;
    }
}

class MedianFinder {
    priority_queue<int> left;
    priority_queue<int, vector<int>, greater<int>> right;
public:
    MedianFinder() {}

    void addNum(int num) {
        left.push(num);
        right.push(left.top()); left.pop();
        if (right.size() > left.size()) {
            left.push(right.top()); right.pop();
        }
    }

    double findMedian() {
        if (left.size() > right.size())
            return left.top();
        return (left.top() + right.top()) / 2.0;
    }
};

type hp struct{ sort.IntSlice }

func (h *hp) Push(x any) { h.IntSlice = append(h.IntSlice, x.(int)) }
func (h *hp) Pop() any {
	a := h.IntSlice
	x := a[len(a)-1]
	h.IntSlice = a[:len(a)-1]
	return x
}

type MedianFinder struct {
	left, right *hp
}

func Constructor() MedianFinder {
	return MedianFinder{&hp{}, &hp{}}
}

func (this *MedianFinder) AddNum(num int) {
	heap.Push(this.left, -num)
	heap.Push(this.right, -heap.Pop(this.left).(int))
	if this.right.Len() > this.left.Len() {
		heap.Push(this.left, -heap.Pop(this.right).(int))
	}
}

func (this *MedianFinder) FindMedian() float64 {
	if this.left.Len() > this.right.Len() {
		return float64(-this.left.IntSlice[0])
	}
	return float64(-this.left.IntSlice[0]+this.right.IntSlice[0]) / 2
}

class MinHeap {
    constructor() { this.a = []; }
    size() { return this.a.length; }
    top() { return this.a[0]; }
    push(x) {
        const a = this.a; a.push(x);
        for (let i = a.length - 1; i > 0;) {
            const p = (i - 1) >> 1;
            if (a[p] <= a[i]) break;
            [a[p], a[i]] = [a[i], a[p]]; i = p;
        }
    }
    pop() {
        const a = this.a, top = a[0], last = a.pop();
        if (a.length) {
            a[0] = last;
            for (let i = 0; ;) {
                let m = i; const l = 2 * i + 1, r = l + 1;
                if (l < a.length && a[l] < a[m]) m = l;
                if (r < a.length && a[r] < a[m]) m = r;
                if (m === i) break;
                [a[m], a[i]] = [a[i], a[m]]; i = m;
            }
        }
        return top;
    }
}

var MedianFinder = function() {
    this.left = new MinHeap();
    this.right = new MinHeap();
};

MedianFinder.prototype.addNum = function(num) {
    this.left.push(-num);
    this.right.push(-this.left.pop());
    if (this.right.size() > this.left.size())
        this.left.push(-this.right.pop());
};

MedianFinder.prototype.findMedian = function() {
    if (this.left.size() > this.right.size())
        return -this.left.top();
    return (-this.left.top() + this.right.top()) / 2;
};

class MinHeap {
    private a: number[] = [];
    size(): number { return this.a.length; }
    top(): number { return this.a[0]; }
    push(x: number): void {
        const a = this.a; a.push(x);
        for (let i = a.length - 1; i > 0;) {
            const p = (i - 1) >> 1;
            if (a[p] <= a[i]) break;
            [a[p], a[i]] = [a[i], a[p]]; i = p;
        }
    }
    pop(): number {
        const a = this.a, top = a[0], last = a.pop()!;
        if (a.length) {
            a[0] = last;
            for (let i = 0; ;) {
                let m = i; const l = 2 * i + 1, r = l + 1;
                if (l < a.length && a[l] < a[m]) m = l;
                if (r < a.length && a[r] < a[m]) m = r;
                if (m === i) break;
                [a[m], a[i]] = [a[i], a[m]]; i = m;
            }
        }
        return top;
    }
}

class MedianFinder {
    private left = new MinHeap();
    private right = new MinHeap();

    constructor() {}

    addNum(num: number): void {
        this.left.push(-num);
        this.right.push(-this.left.pop());
        if (this.right.size() > this.left.size())
            this.left.push(-this.right.pop());
    }

    findMedian(): number {
        if (this.left.size() > this.right.size())
            return -this.left.top();
        return (-this.left.top() + this.right.top()) / 2;
    }
}

载入编辑器，自己跑一遍 →

其他解法

有序数组保存全部顺序，插入代价 O(n)；双堆只维护“较小一半的最大值”和“较大一半的最小值”，不关心内部顺序，插入降到 O(log n)。

解法 2：有序数组 + 二分插入时间 addNum O(n)，findMedian O(1)，空间 O(n)

始终维护一个有序数组。插入时先二分找到位置再插入，取中位数时直接按下标取中间一个或两个元素。

实现简单，还能顺便支持查询任意分位数。插入要搬移元素，最坏 O(n)，数据量大时不如双堆，不过内存搬移很快，实际表现尚可。

from bisect import insort

class MedianFinder:
    def __init__(self):
        self.a = []

    def addNum(self, num: int) -> None:
        insort(self.a, num)

    def findMedian(self) -> float:
        a, n = self.a, len(self.a)
        if n % 2:
            return float(a[n // 2])
        return (a[n // 2 - 1] + a[n // 2]) / 2

class MedianFinder {
    private List<Integer> a = new ArrayList<>();

    public MedianFinder() {}

    public void addNum(int num) {
        int lo = 0, hi = a.size();
        while (lo < hi) {
            int mid = (lo + hi) >>> 1;
            if (a.get(mid) <= num) lo = mid + 1;
            else hi = mid;
        }
        a.add(lo, num);
    }

    public double findMedian() {
        int n = a.size();
        if (n % 2 == 1) return a.get(n / 2);
        return ((double) a.get(n / 2 - 1) + a.get(n / 2)) / 2.0;
    }
}

class MedianFinder {
    vector<int> a;
public:
    MedianFinder() {}

    void addNum(int num) {
        a.insert(upper_bound(a.begin(), a.end(), num), num);
    }

    double findMedian() {
        int n = a.size();
        if (n % 2) return a[n / 2];
        return ((double)a[n / 2 - 1] + a[n / 2]) / 2.0;
    }
};

type MedianFinder struct {
	a []int
}

func Constructor() MedianFinder {
	return MedianFinder{}
}

func (this *MedianFinder) AddNum(num int) {
	i := sort.SearchInts(this.a, num)
	this.a = append(this.a, 0)
	copy(this.a[i+1:], this.a[i:])
	this.a[i] = num
}

func (this *MedianFinder) FindMedian() float64 {
	n := len(this.a)
	if n%2 == 1 {
		return float64(this.a[n/2])
	}
	return float64(this.a[n/2-1]+this.a[n/2]) / 2
}

var MedianFinder = function() {
    this.a = [];
};

MedianFinder.prototype.addNum = function(num) {
    const a = this.a;
    let lo = 0, hi = a.length;
    while (lo < hi) {
        const mid = (lo + hi) >> 1;
        if (a[mid] <= num) lo = mid + 1;
        else hi = mid;
    }
    a.splice(lo, 0, num);
};

MedianFinder.prototype.findMedian = function() {
    const a = this.a, n = a.length;
    if (n % 2) return a[n >> 1];
    return (a[n / 2 - 1] + a[n / 2]) / 2;
};

class MedianFinder {
    private a: number[] = [];

    constructor() {}

    addNum(num: number): void {
        const a = this.a;
        let lo = 0, hi = a.length;
        while (lo < hi) {
            const mid = (lo + hi) >> 1;
            if (a[mid] <= num) lo = mid + 1;
            else hi = mid;
        }
        a.splice(lo, 0, num);
    }

    findMedian(): number {
        const a = this.a, n = a.length;
        if (n % 2) return a[n >> 1];
        return (a[n / 2 - 1] + a[n / 2]) / 2;
    }
}

栈与堆专题的其他题目

查看栈与堆全部题目 →