题目描述
中位数是有序整数列表中的中间值。如果列表的大小是偶数,则没有中间值,中位数是两个中间值的平均值。
例如 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
以内的答案将被接受。
示例 1:
输入
[“MedianFinder”, “addNum”, “addNum”, “findMedian”, “addNum”, “findMedian”]
[[], [1], [2], [], [3], []]
输出
[null, null, null, 1.5, null, 2.0]
解释
MedianFinder medianFinder = new MedianFinder();
medianFinder.addNum(1); // arr = [1]
medianFinder.addNum(2); // arr = [1, 2]
medianFinder.findMedian(); // 返回 1.5 ((1 + 2) / 2)
medianFinder.addNum(3); // arr[1, 2, 3]
medianFinder.findMedian(); // return 2.0
提示:
- -105 <= num <= 105
- 在调用 findMedian 之前,数据结构中至少有一个元素
- 最多 5 * 104 次调用 addNum 和 findMedian
代码及注释
import (
"container/heap"
"sort"
)
// 定义 MedianFinder 结构体
type MedianFinder struct {
queMin, queMax hp // 两个堆
}
// MedianFinder 的构造函数
func Constructor() MedianFinder {
return MedianFinder{}
}
// 添加数字到 MedianFinder 结构体中
func (mf *MedianFinder) AddNum(num int) {
minQ, maxQ := &mf.queMin, &mf.queMax // 定义两个指针,分别指向 queMin 和 queMax
// 如果 minQ 是空的,或者 num 比 minQ 的最小值还小
if minQ.Len() == 0 || num <= -minQ.IntSlice[0] {
heap.Push(minQ, -num) // 将 -num 推入 minQ 堆
// 如果 maxQ 的长度比 minQ 小1,将 minQ 的最小值推入 maxQ
if maxQ.Len()+1 < minQ.Len() {
heap.Push(maxQ, -heap.Pop(minQ).(int))
}
} else {
heap.Push(maxQ, num) // 将 num 推入 maxQ 堆
// 如果 maxQ 的长度大于 minQ,将 maxQ 的最大值推入 minQ
if maxQ.Len() > minQ.Len() {
heap.Push(minQ, -heap.Pop(maxQ).(int))
}
}
}
// 查找中位数
func (mf *MedianFinder) FindMedian() float64 {
minQ, maxQ := mf.queMin, mf.queMax
// 如果 minQ 的长度大于 maxQ,返回 minQ 的最小值
if minQ.Len() > maxQ.Len() {
return float64(-minQ.IntSlice[0])
}
// 否则,返回 maxQ 和 minQ 的中间值
return float64(maxQ.IntSlice[0]-minQ.IntSlice[0]) / 2
}
// 定义 hp 结构体,继承 sort.IntSlice
type hp struct{ sort.IntSlice }
// Push 方法,将元素推入 hp 堆
func (h *hp) Push(v interface{}) {
h.IntSlice = append(h.IntSlice, v.(int))
}
// Pop 方法,从 hp 堆中弹出元素
func (h *hp) Pop() interface{} {
a := h.IntSlice
v := a[len(a)-1]
h.IntSlice = a[:len(a)-1]
return v
}
/**
* Your MedianFinder object will be instantiated and called as such:
* obj := Constructor();
* obj.AddNum(num);
* param_2 := obj.FindMedian();
*/
代码解释
- 结构体定义:
type MedianFinder struct {
queMin, queMax hp
}
定义了一个 MedianFinder
结构体,其中包含两个堆 queMin
和 queMax
。
- 初始化:
func Constructor() MedianFinder {
return MedianFinder{}
}
初始化 MedianFinder
结构体。
- 添加数字到数据结构中:
func (mf *MedianFinder) AddNum(num int) {
minQ, maxQ := &mf.queMin, &mf.queMax
if minQ.Len() == 0 || num <= -minQ.IntSlice[0] {
heap.Push(minQ, -num)
if maxQ.Len()+1 < minQ.Len() {
heap.Push(maxQ, -heap.Pop(minQ).(int))
}
} else {
heap.Push(maxQ, num)
if maxQ.Len() > minQ.Len() {
heap.Push(minQ, -heap.Pop(maxQ).(int))
}
}
}
- 如果
minQ
为空或者num
小于等于minQ
的最大值,则将-num
添加到minQ
中。 - 否则,将
num
添加到maxQ
中。 - 为了保持
minQ
的大小始终大于或等于maxQ
的大小,我们在添加元素后调整两个堆。
- 查找中位数:
func (mf *MedianFinder) FindMedian() float64 {
minQ, maxQ := mf.queMin, mf.queMax
if minQ.Len() > maxQ.Len() {
return float64(-minQ.IntSlice[0])
}
return float64(maxQ.IntSlice[0]-minQ.IntSlice[0]) / 2
}
- 如果
minQ
的大小大于maxQ
,则中位数是-minQ
的最大值。 - 否则,中位数是
(maxQ
的最小值 -minQ
的最大值) / 2。
- 堆的定义:
type hp struct{ sort.IntSlice }
func (h *hp) Push(v interface{}) { h.IntSlice = append(h.IntSlice, v.(int)) }
func (h *hp) Pop() interface{} { a := h.IntSlice; v := a[len(a)-1]; h.IntSlice = a[:len(a)-1]; return v }
定义了一个堆 hp
,这个堆基于 sort.IntSlice
实现,同时提供了 Push
和 Pop
方法。