u
This commit is contained in:
@@ -0,0 +1,134 @@
|
||||
package top100liked
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// https://leetcode.cn/problems/find-median-from-data-stream/?envType=study-plan-v2&envId=top-100-liked
|
||||
|
||||
type heap[T any] struct {
|
||||
data []T
|
||||
compareFunc func(T, T) bool
|
||||
}
|
||||
|
||||
func NewHeap[T any](compareFunc func(T, T) bool) heap[T] {
|
||||
return heap[T]{
|
||||
compareFunc: compareFunc,
|
||||
}
|
||||
}
|
||||
|
||||
func (h *heap[T]) Up(idx int) {
|
||||
for (idx-1)/2 >= 0 {
|
||||
cur := h.data[idx]
|
||||
p := h.data[(idx-1)/2]
|
||||
if h.compareFunc(cur, p) {
|
||||
h.data[idx], h.data[(idx-1)/2] = h.data[(idx-1)/2], h.data[idx]
|
||||
idx = (idx - 1) / 2
|
||||
} else {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (h *heap[T]) Down(idx int) {
|
||||
for {
|
||||
l, r := 2*idx+1, 2*idx+2
|
||||
m := idx
|
||||
if l < len(h.data) && h.compareFunc(h.data[l], h.data[m]) {
|
||||
m = l
|
||||
}
|
||||
if r < len(h.data) && h.compareFunc(h.data[r], h.data[m]) {
|
||||
m = r
|
||||
}
|
||||
if m == idx {
|
||||
break
|
||||
}
|
||||
h.data[idx], h.data[m] = h.data[m], h.data[idx]
|
||||
idx = m
|
||||
}
|
||||
}
|
||||
|
||||
func (h *heap[T]) Push(val T) {
|
||||
h.data = append(h.data, val)
|
||||
h.Up(len(h.data) - 1)
|
||||
}
|
||||
|
||||
func (h *heap[T]) Pop() T {
|
||||
if len(h.data) <= 0 {
|
||||
var zero T
|
||||
return zero
|
||||
}
|
||||
old := h.data[0]
|
||||
h.data[0] = h.data[len(h.data)-1]
|
||||
h.data = h.data[:len(h.data)-1]
|
||||
h.Down(0)
|
||||
return old
|
||||
}
|
||||
|
||||
func (h *heap[T]) Top() T {
|
||||
if len(h.data) > 0 {
|
||||
return h.data[0]
|
||||
}
|
||||
var zero T
|
||||
return zero
|
||||
}
|
||||
|
||||
func (h *heap[T]) Len() int {
|
||||
return len(h.data)
|
||||
}
|
||||
|
||||
type MedianFinder struct {
|
||||
left heap[int] // 小
|
||||
right heap[int] // 大
|
||||
}
|
||||
|
||||
func Constructor() MedianFinder {
|
||||
return MedianFinder{
|
||||
left: NewHeap(func(i1, i2 int) bool { return i1 > i2 }),
|
||||
right: NewHeap(func(i1, i2 int) bool { return i1 < i2 }),
|
||||
}
|
||||
}
|
||||
|
||||
func (this *MedianFinder) AddNum(num int) {
|
||||
if this.left.Len() == 0 {
|
||||
this.left.Push(num)
|
||||
return
|
||||
}
|
||||
if num >= this.left.Top() {
|
||||
this.right.Push(num)
|
||||
} else {
|
||||
this.left.Push(num)
|
||||
}
|
||||
if this.left.Len() > this.right.Len()+1 {
|
||||
v := this.left.Pop()
|
||||
this.right.Push(v)
|
||||
}
|
||||
if this.right.Len() > this.left.Len()+1 {
|
||||
v := this.right.Pop()
|
||||
this.left.Push(v)
|
||||
}
|
||||
}
|
||||
|
||||
func (this *MedianFinder) FindMedian() float64 {
|
||||
switch (this.left.Len() + this.right.Len()) % 2 {
|
||||
case 0:
|
||||
return (float64(this.left.Top()) + float64(this.right.Top())) / 2
|
||||
case 1:
|
||||
if this.left.Len() > this.right.Len() {
|
||||
return float64(this.left.Top())
|
||||
} else {
|
||||
return float64(this.right.Top())
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func Test(t *testing.T) {
|
||||
medianFinder := Constructor()
|
||||
medianFinder.AddNum(1) // arr = [1]
|
||||
medianFinder.AddNum(2) // arr = [1, 2]
|
||||
fmt.Println(medianFinder.FindMedian()) // 返回 1.5 ((1 + 2) / 2)
|
||||
medianFinder.AddNum(3) // arr[1, 2, 3]
|
||||
fmt.Println(medianFinder.FindMedian()) // return 2.0
|
||||
}
|
||||
Reference in New Issue
Block a user