【路飞】算法与数据结构-数据流的中位数

282 阅读2分钟

不管全世界所有人怎么说,我都认为自己的感受才是正确的。无论别人怎么看,我绝不打乱自己的节奏。喜欢的事自然可以坚持,不喜欢的怎么也长久不了。

LeetCode:原题地址

题目要求

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

例如,

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

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

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

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

示例 1:

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

进阶:

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

思路

此题需要两个堆,大顶堆维护前一半的数据,小顶堆维护后一半的数据; 如果数据个数为n,如果n是偶数,则大顶堆和小顶堆维护一样的个数;如果n为奇数,则大顶堆的个数比小顶堆的个数多1; 插入新数据时,判断它与大顶堆堆顶的大小,如果小于堆顶元素,则插入大顶堆,否则插入小顶堆; 插入后,两个堆的个数不满足要求的话,就根据要求,从数量过多的堆删除堆顶元素做堆化调整,将堆顶元素插入到另一个堆的堆底,做堆化调整。

代码

/**
 * initialize your data structure here.
 */
var MedianFinder = function () {
	this.maxHeap = new Heap()
	this.minHeap = new Heap()
}

/**
 * @param {number} num
 * @return {void}
 */
MedianFinder.prototype.addNum = function (num) {
	let maxSize = this.maxHeap.getSize()
	let minSize = this.minHeap.getSize()
	if (maxSize === 0) {
		this.maxHeap.insert(num, false)
	} else {
		if (this.maxHeap.getTop() > num) {
			this.maxHeap.insert(num, false)
		} else {
			this.minHeap.insert(num, true)
		}
		maxSize = this.maxHeap.getSize()
		minSize = this.minHeap.getSize()		
		if (minSize > maxSize) {
			// 需要从小顶堆删除堆顶元素,并放入到大顶堆
			let top = this.minHeap.pop(true)
			this.maxHeap.insert(top, false)
		} else if (maxSize - 1 > minSize) {
			// 需要从大顶堆删除堆顶元素,并放入到小顶堆
			let top = this.maxHeap.pop(false)
			this.minHeap.insert(top, true)
		}
	}
}

/**
 * @return {number}
 */
MedianFinder.prototype.findMedian = function () {
	let maxSize = this.maxHeap.getSize()
	let minSize = this.minHeap.getSize()
	if (maxSize === minSize) {
		return (this.maxHeap.getTop() + this.minHeap.getTop()) / 2
	} else {
		return this.maxHeap.getTop()
	}
}

class Heap {
	constructor() {
		this.data = [0]
		this.count = 0
	}

	getTop() {
		return this.data[1]
	}

	pop(minFlag) {
		let top = this.data[1]
		if (this.count > 1) {
			this.data[1] = this.data[this.count]
			this.data.splice(this.count, 1)
			this.count--
			let i = 1
			if (minFlag) {
				while (true) {
					let maxPos = i
					if (
						i * 2 <= this.count &&
						this.data[i] > this.data[i * 2]
					) {
						maxPos = i * 2
					}
					if (
						i * 2 + 1 <= this.count &&
						this.data[maxPos] > this.data[i * 2 + 1]
					) {
						maxPos = i * 2 + 1
					}
					if (i === maxPos) {
						break
					}
					this.swap(i, maxPos)
					i = maxPos
				}
			} else {
				while (true) {
					let maxPos = i
					if (
						i * 2 <= this.count &&
						this.data[i] < this.data[i * 2]
					) {
						maxPos = i * 2
					}
					if (
						i * 2 + 1 <= this.count &&
						this.data[maxPos] < this.data[i * 2 + 1]
					) {
						maxPos = i * 2 + 1
					}
					if (i === maxPos) {
						break
					}
					this.swap(i, maxPos)
					i = maxPos
				}
			}
		} else {
			this.data.splice(this.count, 1)
			this.count--
		}
		return top
	}

	insert(val, minFlag) {
		this.count++
		this.data[this.count] = val
		let i = this.count
		if (minFlag) {
			while (
				parseInt(i / 2) > 0 &&
				this.data[i] < this.data[parseInt(i / 2)]
			) {
				this.swap(i, parseInt(i / 2))
				i = parseInt(i / 2)
			}
		} else {
			while (
				parseInt(i / 2) > 0 &&
				this.data[i] > this.data[parseInt(i / 2)]
			) {
				this.swap(i, parseInt(i / 2))
				i = parseInt(i / 2)
			}
		}
	}

	getSize() {
		return this.count
	}

	swap(i, j) {
		let temp = this.data[i]
		this.data[i] = this.data[j]
		this.data[j] = temp
	}
}

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