leetcode 295. 数据流的中位数leetcode 剑指 Offer 41. 数据流中的中位数。 NC131 随

思路

利用两个堆，将当前从数据流中得到的数据分成两部分。两部分数据个数相差至多1。一部分放在一个大顶堆A里，另一部分放在一个小顶堆B里。将数据流中的数加入堆时，保证A中所有数都小于或等于B中的数。具体细节见方法func (this *MedianFinder) AddNum(num int)。这样，计算中位数就简单了。

下面给出力扣上题目的解答(已提交通过)。牛客题目的解法在此基础上稍加修改即可。

type MedianFinder struct {
	maxHeap *MaxHeap
	minHeap *MinHeap
}


/** initialize your data structure here. */
func Constructor() MedianFinder {
	return MedianFinder{
		maxHeap: NewMaxHeap(),
		minHeap: NewMinHeap(),
	}
}


func (this *MedianFinder) AddNum(num int)  {
	if this.maxHeap.IsEmpty() { // 优先放入大顶堆
		this.maxHeap.Push(num)
	} else if num <= this.maxHeap.Top() { // 优先放入大顶堆
		this.maxHeap.Push(num)
	} else {
		this.minHeap.Push(num)
	}
	total := this.maxHeap.Len() + this.minHeap.Len()
	half := total/2
        // total&1 == 0时，total是偶数，每个堆刚好分一半数据。
        // total&1 != 0时，total&1结果是1，total是奇数,
        // 让大顶堆中的数据个数比小顶堆中的多1，简化获取中位数时的逻辑(请看函数FindMedian)。
	if this.maxHeap.Len() > half+total&1 { // 个数多了就调整
		this.minHeap.Push(this.maxHeap.Top())
		this.maxHeap.Pop()
	}
        if this.minHeap.Len() > half { // 个数多了就调整
		this.maxHeap.Push(this.minHeap.Top())
		this.minHeap.Pop()
	}
}


func (this *MedianFinder) FindMedian() float64 {
        // 两个堆中的数据量相同
	if this.maxHeap.Len() == this.minHeap.Len() {
		return (float64(this.maxHeap.Top())+float64(this.minHeap.Top()))/2
	}
        // 数据量不同，直接取大顶堆的堆顶
	return float64(this.maxHeap.Top())
}


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

// 将大顶堆和小顶堆的公共部分抽取出来，实现为Heap
type Heap struct {
	data []int
	size int
        // 父节点和子节点是否“正序”。
        // true-“正序”。
        // 由大顶堆/小顶堆的特性决定。
	isParentAndChildOrdered func(parentVal, childVal int) bool
        // 在兄弟节点比较时，判断是否选择“靠右”(索引值更大)的兄弟节点。
        // true-选择“靠右”节点。
        // 由大顶堆/小顶堆的特性决定。
	shouldChooseRightSibling func(leftVal, rightVal int) bool
}

func NewHeap(isParentAndChildOrdered func(parentVal, childVal int) bool,
	shouldChooseRightSibling func(leftVal, rightVal int) bool) *Heap {
	return &Heap{
		isParentAndChildOrdered: isParentAndChildOrdered,
		shouldChooseRightSibling: shouldChooseRightSibling,
	}
}

func (o *Heap) Len() int {
	return o.size
}

func (o *Heap) IsEmpty() bool {
	return o.size == 0
}

func (o *Heap) Top() int {
	if o.IsEmpty() {
		panic("Heap is empty")
	}
	return o.data[0]
}

func (o *Heap) Push(n int) {
	if o.size < len(o.data) {
		o.data[o.size] = n
	} else {
		o.data = append(o.data, n)
	}
	o.size++
	o.siftUp()
}

func (o *Heap) siftUp() {
	if o.size < 2 {
		return
	}
	var (
		child = o.size-1
		childV = o.data[child]
	)
	for child > 0 {
		p := (child-1)>>1
		pV := o.data[p]
		if o.isParentAndChildOrdered(pV, childV) {
			break
		}
		o.data[child] = pV
		child = p
	}
	o.data[child] = childV
}

func (o *Heap) Pop() {
	if o.IsEmpty() {
		panic("Heap is empty")
	}
	o.data[0] = o.data[o.size-1]
	o.size--
	o.siftDown()
}

func (o *Heap) siftDown() {
	if o.size < 2 {
		return
	}
	var (
		p int
		pV = o.data[p]
	)
	for {
		child := 2*p + 1
		if child >= o.size {
			break
		}
		if child+1 < o.size && 
                   o.shouldChooseRightSibling(o.data[child], o.data[child+1]) {
			child++
		}
		childV := o.data[child]
		if o.isParentAndChildOrdered(pV, childV) {
			break
		}
		o.data[p] = childV
		p = child
	}
	o.data[p] = pV
}

// 大顶堆
type MaxHeap struct {
	*Heap
}

func NewMaxHeap() *MaxHeap {
	o := &MaxHeap{}
	o.Heap = NewHeap(o.isParentAndChildOrdered, o.shouldChooseRightSibling)
	return o
}

func (MaxHeap) isParentAndChildOrdered(parentVal, childVal int) bool {
	return parentVal >= childVal
}

func (MaxHeap) shouldChooseRightSibling(leftVal, rightVal int) bool {
	return leftVal < rightVal
}

// 小顶堆
type MinHeap struct {
	*Heap
}

func NewMinHeap() *MinHeap {
	o := &MinHeap{}
	o.Heap = NewHeap(o.isParentAndChildOrdered, o.shouldChooseRightSibling)
	return o
}

func (MinHeap) isParentAndChildOrdered(parentVal, childVal int) bool {
	return parentVal <= childVal
}

func (MinHeap) shouldChooseRightSibling(leftVal, rightVal int) bool {
	return leftVal > rightVal
}

两个堆，持有所有已知数据，空间复杂度是 $O(n)$ 。
方法func (this *MedianFinder) AddNum(num int)，时间复杂度 $O(logn)$ 。
方法func (this *MedianFinder) FindMedian() float64，时间复杂度 $O(1)$ 。

上面的代码是“嵌入(继承)+函数指针”，下面的实现改为“嵌入(组合)+接口”(已提交通过)。

type MedianFinder struct {
	maxHeap *Heap
	minHeap *Heap
}


/** initialize your data structure here. */
func Constructor() MedianFinder {
	return MedianFinder{
		maxHeap: NewHeap(&MaxHeapCompare{}),
		minHeap: NewHeap(&MinHeapCompare{}),
	}
}


func (this *MedianFinder) AddNum(num int)  {
	if this.maxHeap.IsEmpty() {
		this.maxHeap.Push(num)
	} else if num <= this.maxHeap.Top() {
		this.maxHeap.Push(num)
	} else {
		this.minHeap.Push(num)
	}

	total := this.maxHeap.Len() + this.minHeap.Len()
	half := total/2
	if total&1 == 0 { // 偶数
		for this.maxHeap.Len() > half {
			this.minHeap.Push(this.maxHeap.Top())
			this.maxHeap.Pop()
		}
		for this.minHeap.Len() > half {
			this.maxHeap.Push(this.minHeap.Top())
			this.minHeap.Pop()
		}
	} else {
		oneMore := half+1
		for this.maxHeap.Len() > oneMore {
			this.minHeap.Push(this.maxHeap.Top())
			this.maxHeap.Pop()
		}
		for this.minHeap.Len() > half {
			this.maxHeap.Push(this.minHeap.Top())
			this.minHeap.Pop()
		}
	}
}


func (this *MedianFinder) FindMedian() float64 {
	if this.maxHeap.Len() == this.minHeap.Len() {
		return (float64(this.maxHeap.Top())+float64(this.minHeap.Top()))/2
	}
	return float64(this.maxHeap.Top())
}


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

type HeapCompare interface {
	IsParentAndChildOrdered(parentVal, childVal int) bool
	ShouldChooseRightSibling(leftVal, rightVal int) bool
}

type Heap struct {
	data []int
	size int
	HeapCompare
}

func NewHeap(cmp HeapCompare) *Heap {
	return &Heap{
		HeapCompare: cmp,
	}
}

func (o *Heap) Len() int {
	return o.size
}

func (o *Heap) IsEmpty() bool {
	return o.size == 0
}

func (o *Heap) Top() int {
	if o.IsEmpty() {
		panic("Heap is empty")
	}
	return o.data[0]
}

func (o *Heap) Push(n int) {
	if o.size < len(o.data) {
		o.data[o.size] = n
	} else {
		o.data = append(o.data, n)
	}
	o.size++
	o.siftUp()
}

func (o *Heap) siftUp() {
	if o.size < 2 {
		return
	}
	var (
		child = o.size-1
		childV = o.data[child]
	)
	for child > 0 {
		p := (child-1)>>1
		pV := o.data[p]
		if o.IsParentAndChildOrdered(pV, childV) {
			break
		}
		o.data[child] = pV
		child = p
	}
	o.data[child] = childV
}

func (o *Heap) Pop() {
	if o.IsEmpty() {
		panic("Heap is empty")
	}
	o.data[0] = o.data[o.size-1]
	o.size--
	o.siftDown()
}

func (o *Heap) siftDown() {
	if o.size < 2 {
		return
	}
	var (
		p int
		pV = o.data[p]
	)
	for {
		child := 2*p + 1
		if child >= o.size {
			break
		}
		if child+1 < o.size && 
                   o.ShouldChooseRightSibling(o.data[child], o.data[child+1]) {
			child++
		}
		childV := o.data[child]
		if o.IsParentAndChildOrdered(pV, childV) {
			break
		}
		o.data[p] = childV
		p = child
	}
	o.data[p] = pV
}

type MaxHeapCompare struct {}

func (o *MaxHeapCompare) IsParentAndChildOrdered(parentVal, childVal int) bool {
	return parentVal >= childVal
}

func (o *MaxHeapCompare) ShouldChooseRightSibling(leftVal, rightVal int) bool {
	return leftVal < rightVal
}

type MinHeapCompare struct {}

func (o *MinHeapCompare) IsParentAndChildOrdered(parentVal, childVal int) bool {
	return parentVal <= childVal
}

func (o *MinHeapCompare) ShouldChooseRightSibling(leftVal, rightVal int) bool {
	return leftVal > rightVal
}