41.1 数据流中的中位数

题目链接

牛客网 (opens new window)

力扣 (opens new window)

题目描述

如何得到一个数据流中的中位数？如果从数据流中读出奇数个数值，那么中位数就是所有数值排序之后位于中间的数值。如果从数据流中读出偶数个数值，那么中位数就是所有数值排序之后中间两个数的平均值。

解题思路

type MedianFinder struct {
	minHeap []int
	maxHeap []int
	length  int
}

/** initialize your data structure here. */
func Constructor() MedianFinder {
	return MedianFinder{
		minHeap: []int{},
		maxHeap: []int{},
		length:  0, // 记录存入元素的个数
	}
}

// 上浮 用于堆尾插入元素 调整堆  cmp函数是为了控制是大顶堆还是小顶堆
func upAdjust(arr []int, cmp func(x, y int) bool) {
	if len(arr) == 0 {
		return
	}
	length := len(arr)
	child := length - 1
	parent := (child - 1) / 2
	newValue := arr[child]
	for child > 0 {
		if cmp(newValue, arr[parent]) { // 满足条件上浮  小顶堆 就是 newValue 小于  交换 当前 大顶堆就是大于当前
			arr[child] = arr[parent]
			child = parent
			parent = (child - 1) / 2
		} else {
			break
		}
	}
	arr[child] = newValue
}

// 下沉 用于删除堆顶元素后 调整堆
func downAdjust(arr []int, i int, cmp func(x, y int) bool) {
	if len(arr) == 0 {
		return
	}
	parent := i
	child := 2*parent + 1
	rootValue := arr[parent]
	for child < len(arr) {
		if child+1 < len(arr) && cmp(arr[child+1], arr[child]) { // 小顶堆 就是寻找小的 child 用于以后比较 大顶堆就是寻找大的
			child++
		}
		if cmp(arr[child], rootValue) { // 满足条件下沉  小顶堆 就是 rootValue 大于当前 交换   大顶堆就是 小于当前
			arr[parent] = arr[child]
			parent = child
			child = parent*2 + 1
		} else {
			break
		}
	}
	arr[parent] = rootValue
}
// 维护小顶堆
func cmpMin(x, y int) bool {
	return x < y
}
// 维护大顶堆
func cmpMax(x, y int) bool {
	return x > y
}

// 添加数
func (this *MedianFinder) AddNum(num int) {
	this.length++
	if this.length%2 == 0 {
		//fmt.Println("min:", this.minHeap)
		// 存入小顶堆 并调整
		this.minHeap = append(this.minHeap, num)
		upAdjust(this.minHeap, cmpMin)
		// 弹出小顶堆栈顶存入大顶堆
		this.maxHeap = append(this.maxHeap, this.minHeap[0])
		Pop(&this.minHeap)
		// 调整弹出元素后的小顶堆 和 新加元素的大顶堆
		downAdjust(this.minHeap, 0, cmpMin)
		upAdjust(this.maxHeap, cmpMax)
	} else if this.length%2 == 1 { // 保证堆平衡 维护大小顶堆长度 相差不大于1
		//fmt.Println("max:", this.maxHeap)
		// 存入大顶堆 并调整
		this.maxHeap = append(this.maxHeap, num)
		upAdjust(this.maxHeap, cmpMax)
		// 弹出大顶堆堆顶 并存入小顶堆
		this.minHeap = append(this.minHeap, this.maxHeap[0])
		Pop(&this.maxHeap)
		//fmt.Println("max:", this.maxHeap)
		// 调整弹出元素后的大顶堆 和 新加元素的小顶堆
		downAdjust(this.maxHeap, 0, cmpMax)
		upAdjust(this.minHeap, cmpMin)
		//fmt.Println("min:", this.minHeap)
	}
	return
}

// 弹出堆顶元素 并将堆尾元素放入堆顶  调整堆的长度
func Pop(arr *[]int) { // 这块坑我好久  对于切片的更改  只更改值可以使用值形式 但是更改长度 需要使用指针  很重要
	length := len(*arr)
	if length == 1 { // 若堆只有一个元素 那么无法交换 返回空数组
		*arr = (*arr)[:length-1]
	}
	if length > 1 {
		(*arr)[0] = (*arr)[length-1]
		*arr = (*arr)[:length-1]
	}
}

// 计算中位数
func (this *MedianFinder) FindMedian() float64 {
	if this.length%2 == 1 {
		return float64(this.minHeap[0])
	}
	if this.length%2 == 0 {
		return (float64(this.maxHeap[0]) + float64(this.minHeap[0])) / 2
	}
	return 0
}


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