mrsen/stl-go
1
0
Fork 0
mirror of https://gitee.com/NorthCityChen/stl-go.git synced 2025-05-25 20:11:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: heap

Browse Source
This commit is contained in:
NorthCityChen 2022-04-26 15:57:00 +08:00
parent 0ab851e49f
commit e789b95c8f
2 changed files with 36 additions and 43 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
app.go
View File

@ -1,6 +1,6 @@
/* /*
* @Author: NorthCity1984 * @Author: NorthCity1984
* @LastEditTime: 2022-04-26 14:25:09 * @LastEditTime: 2022-04-26 15:55:37
* @Description: * @Description:
* @Website: https://grimoire.cn * @Website: https://grimoire.cn
* Copyright (c) NorthCity1984 All rights reserved. * Copyright (c) NorthCity1984 All rights reserved.
@ -13,14 +13,18 @@ import (
"gitee.com/NorthCityChen/stl-go/heap" "gitee.com/NorthCityChen/stl-go/heap"
) )
// import "container/heap"
var a = []int{1, 2, 65, 4, 5} var a = []int{1, 2, 65, 4, 5}
func main() { func main() {
h := heap.Init(a, false) // <: 小根堆 >: 大根堆
f := func(t1, t2 int) bool { return t1 > t2 }
h := heap.Init(a, f)
fmt.Println(h) fmt.Println(h)
h.Push(77)
fmt.Println(h) fmt.Println(h)
fmt.Println(h.Top())
fmt.Println(h.Pop()) fmt.Println(h.Pop())
fmt.Println(h.Pop()) fmt.Println(h.Pop())
fmt.Println(h.Pop()) fmt.Println(h.Pop())

67
heap/heap.go
View File

@ -1,6 +1,6 @@
/* /*
* @Author: NorthCity1984 * @Author: NorthCity1984
* @LastEditTime: 2022-04-26 14:24:36 * @LastEditTime: 2022-04-26 15:55:32
* @Description: * @Description:
* @Website: https://grimoire.cn * @Website: https://grimoire.cn
* Copyright (c) NorthCity1984 All rights reserved. * Copyright (c) NorthCity1984 All rights reserved.
@ -14,52 +14,37 @@ type Number interface {
type Heap[T Number] struct { type Heap[T Number] struct {
array []T array []T
len int len int
isMax bool cmp func(T, T) bool
} }
func (h *Heap[T]) up(x int) { func (h *Heap[T]) up(x int) {
if h.isMax { // if h.isMax {
for x > 1 && h.array[x] > h.array[x/2] {
h.array[x], h.array[x/2] = h.array[x/2], h.array[x] for x > 1 && h.cmp(h.array[x], h.array[x/2]) {
x /= 2 h.array[x], h.array[x/2] = h.array[x/2], h.array[x]
} x /= 2
} else {
for x > 1 && h.array[x] < h.array[x/2] {
h.array[x], h.array[x/2] = h.array[x/2], h.array[x]
x /= 2
}
} }
// } else {
// for x > 1 && h.array[x] < h.array[x/2] {
// h.array[x], h.array[x/2] = h.array[x/2], h.array[x]
// x /= 2
// }
// }
} }
func (h *Heap[T]) down(x int) { func (h *Heap[T]) down(x int) {
// fmt.Println(h.len) for x*2 <= h.len {
if h.isMax { t := x * 2
for x*2 <= h.len { if t+1 <= h.len && h.cmp(h.array[t+1], h.array[t]) {
t := x * 2 t++
if t+1 <= h.len && h.array[t+1] > h.array[t] {
t++
}
if h.array[t] <= h.array[x] {
break
}
h.array[x], h.array[t] = h.array[t], h.array[x]
x = t
} }
} else { if !h.cmp(h.array[t], h.array[x]) {
for x*2 <= h.len { break
t := x * 2
if t+1 <= h.len && h.array[t+1] > h.array[t] {
t++
}
if h.array[t] >= h.array[x] {
break
}
h.array[x], h.array[t] = h.array[t], h.array[x]
x = t
} }
h.array[x], h.array[t] = h.array[t], h.array[x]
x = t
} }
} }
func (h *Heap[T]) Push(v T) { func (h *Heap[T]) Push(v T) {
@ -76,6 +61,7 @@ func (h *Heap[T]) Pop() T {
h.array[1], h.array[h.len] = h.array[h.len], h.array[1] h.array[1], h.array[h.len] = h.array[h.len], h.array[1]
h.array = h.array[:h.len] h.array = h.array[:h.len]
h.len-- h.len--
// fmt.Println(h)
h.down(1) h.down(1)
return ret return ret
@ -85,14 +71,17 @@ func (h *Heap[T]) Top() T {
return h.array[1] return h.array[1]
} }
func Init[T Number](array []T, isMax bool) Heap[T] { func Init[T Number](array []T, cmp func(T, T) bool) Heap[T] {
h := Heap[T]{ h := Heap[T]{
array: append([]T{0}, array...), array: append([]T{0}, array...),
len: len(array), len: len(array),
isMax: isMax, cmp: cmp,
} }
for i := len(h.array) - 1; i >= 1; i-- { for i := h.len; i >= 1; i-- {
h.down(i) h.down(i)
} }
// for i := 1; i <= h.len; i++ {
// h.up(i)
// }
return h return h
} }