295. 数据流的中位数

🔴   🔖  设计 双指针 数据流 排序 堆（优先队列）  🔗 LeetCode

题目

The median is the middle value in an ordered integer list. If the size of the list is even, there is no middle value, and the median is the mean of the two middle values.

• For example, for arr = [2,3,4], the median is 3.
• For example, for arr = [2,3], the median is (2 + 3) / 2 = 2.5.

Implement the MedianFinder class:

• MedianFinder() initializes the MedianFinder object.
• void addNum(int num) adds the integer num from the data stream to the data structure.
• double findMedian() returns the median of all elements so far. Answers within 10-5 of the actual answer will be accepted.

Example 1:

Input

["MedianFinder", "addNum", "addNum", "findMedian", "addNum", "findMedian"]

[[], [1], [2], [], [3], []]

Output

[null, null, null, 1.5, null, 2.0]

Explanation

MedianFinder medianFinder = new MedianFinder();

medianFinder.addNum(1); // arr = [1]

medianFinder.addNum(2); // arr = [1, 2]

medianFinder.findMedian(); // return 1.5 (i.e., (1 + 2) / 2)

medianFinder.addNum(3); // arr[1, 2, 3]

medianFinder.findMedian(); // return 2.0

Constraints:

• -10^5 <= num <= 10^5
• There will be at least one element in the data structure before calling findMedian.
• At most 5 * 10^4 calls will be made to addNum and findMedian.

Follow up:

• If all integer numbers from the stream are in the range [0, 100], how would you optimize your solution?
• If 99% of all integer numbers from the stream are in the range [0, 100], how would you optimize your solution?

题目大意

中位数是有序整数列表中的中间值。如果列表的大小是偶数，则没有中间值，中位数是两个中间值的平均值。

• 例如 arr = [2,3,4] 的中位数是 3 。
• 例如 arr = [2,3] 的中位数是 (2 + 3) / 2 = 2.5 。

实现 MedianFinder 类:

• MedianFinder() 初始化 MedianFinder 对象。

• void addNum(int num) 将数据流中的整数 num 添加到数据结构中。

• double findMedian() 返回到目前为止所有元素的中位数。与实际答案相差 10^-5 以内的答案将被接受。

解题思路

可以使用两个堆来解决问题。

• 初始化一个小顶堆 small 和一个大顶堆 large 来存储数据；

• 求中位数：

  • 当两个堆的长度一样时，两个堆顶的平均数就是中位数；
  • 当两个堆的长度不一样时，更长的那个堆的堆顶就是中位数；

• 添加数据：

  • 如果小顶堆 small 的数据比大顶堆 large 的数据多，那么将数据添加到 small 中，再将 small 的堆顶（也即最小值）推出，推入到 large 中，如此便可以保证 small 中的数据都大于 large 中的数；
  • 反之，如果小顶堆 small 的数据比大顶堆 large 的数据少，那么将数据添加到 large 中，再将 large 的堆顶（也即最大值）推出，推入到 small 中，如此便可以保证 small 中的数据都大于 large 中的数；

• 时间复杂度：

  • addNum: O(logn)，其中 n 为累计添加的数的数量。
  • findMedian: O(1)。

• 空间复杂度： O(n)，主要为优先队列的开销。

代码

var MedianFinder = function () {
	// 小顶堆
	this.small = new MinPriorityQueue();
	// 大顶堆
	this.large = new MaxPriorityQueue();
};

/**
 * @param {number} num
 * @return {void}
 */
MedianFinder.prototype.addNum = function (num) {
	if (this.small.size() >= this.large.size()) {
		this.small.insert(num);
		this.large.insert(this.small.pop());
	} else {
		this.large.insert(num);
		this.small.insert(this.large.pop());
	}
};

/**
 * @return {number}
 */
MedianFinder.prototype.findMedian = function () {
	const lenSmall = this.small.size(),
		lenLarge = this.large.size();

	// 如果元素不一样多，多的那个堆的堆顶元素就是中位数
	if (lenSmall > lenLarge) {
		return this.small.heap[0];
	} else if (lenSmall < lenLarge) {
		return this.large.heap[0];
	}

	// 如果元素一样多，两个堆堆顶元素的平均数是中位数
	return (this.small.heap[0] + this.large.heap[0]) / 2;
};

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

class MaxPriorityQueue {
	constructor() {
		this.heap = [];
	}
	size() {
		return this.heap.length;
	}
	insert(value) {
		this.heap.push(value);
		this.heapifyUp(this.heap.length - 1);
	}
	pop() {
		if (this.heap.length == 0) return null;
		const first = this.heap[0],
			last = this.heap.pop();
		if (this.heap.length > 0) {
			this.heap[0] = last;
			this.heapifyDown(0);
		}
		return first;
	}
	heapifyUp(i) {
		while (i) {
			const parent = ((i - 1) / 2) | 0;
			if (this.heap[i] > this.heap[parent]) {
				[this.heap[i], this.heap[parent]] = [this.heap[parent], this.heap[i]];
				i = parent;
			} else {
				break;
			}
		}
	}
	heapifyDown(i) {
		let left = i * 2 + 1,
			right = i * 2 + 2,
			min = i;
		if (left < this.heap.length && this.heap[min] < this.heap[left]) {
			min = left;
		}

		if (right < this.heap.length && this.heap[min] < this.heap[right]) {
			min = right;
		}

		if (min !== i) {
			[this.heap[i], this.heap[min]] = [this.heap[min], this.heap[i]];
			this.heapifyDown(min);
		}
	}
}

class MinPriorityQueue {
	constructor() {
		this.heap = [];
	}
	size() {
		return this.heap.length;
	}
	insert(value) {
		this.heap.push(value);
		this.heapifyUp(this.heap.length - 1);
	}
	pop() {
		if (this.heap.length == 0) return null;
		const first = this.heap[0],
			last = this.heap.pop();
		if (this.heap.length > 0) {
			this.heap[0] = last;
			this.heapifyDown(0);
		}
		return first;
	}
	heapifyUp(i) {
		while (i) {
			const parent = ((i - 1) / 2) | 0;
			if (this.heap[i] < this.heap[parent]) {
				[this.heap[i], this.heap[parent]] = [this.heap[parent], this.heap[i]];
				i = parent;
			} else {
				break;
			}
		}
	}
	heapifyDown(i) {
		let left = i * 2 + 1,
			right = i * 2 + 2,
			min = i;
		if (left < this.heap.length && this.heap[min] > this.heap[left]) {
			min = left;
		}

		if (right < this.heap.length && this.heap[min] > this.heap[right]) {
			min = right;
		}

		if (min !== i) {
			[this.heap[i], this.heap[min]] = [this.heap[min], this.heap[i]];
			this.heapifyDown(min);
		}
	}
}

相关题目

题号	标题	题解	标签	难度
480	滑动窗口中位数		数组 哈希表 滑动窗口 1+
1825	求出 MK 平均值		设计 队列 数据流 2+
2102	序列顺序查询		设计 数据流 有序集合 1+
3107	使数组中位数等于 K 的最少操作数		贪心 数组 排序