源码
package heap
import "sort"
// The Interface type describes the requirements
// for a type using the routines in this package.
// Any type that implements it may be used as a
// min-heap with the following invariants (established after
// Init has been called or if the data is empty or sorted):
//
// !h.Less(j, i) for 0 <= i < h.Len() and 2*i+1 <= j <= 2*i+2 and j < h.Len()
//
// Note that Push and Pop in this interface are for package heap's
// implementation to call. To add and remove things from the heap,
// use heap.Push and heap.Pop.
type Interface interface {
sort.Interface
Push(x interface{}) // add x as element Len()
Pop() interface{} // remove and return element Len() - 1.
}
// Init establishes the heap invariants required by the other routines in this package.
// Init is idempotent with respect to the heap invariants
// and may be called whenever the heap invariants may have been invalidated.
// The complexity is O(n) where n = h.Len().
func Init(h Interface) {
// heapify
n := h.Len()
for i := n/2 - 1; i >= 0; i-- {
down(h, i, n)
}
}
// Push pushes the element x onto the heap.
// The complexity is O(log n) where n = h.Len().
func Push(h Interface, x interface{}) {
h.Push(x)
up(h, h.Len()-1)
}
// Pop removes and returns the minimum element (according to Less) from the heap.
// The complexity is O(log n) where n = h.Len().
// Pop is equivalent to Remove(h, 0).
func Pop(h Interface) interface{} {
n := h.Len() - 1
h.Swap(0, n)
down(h, 0, n)
return h.Pop()
}
// Remove removes and returns the element at index i from the heap.
// The complexity is O(log n) where n = h.Len().
func Remove(h Interface, i int) interface{} {
n := h.Len() - 1
if n != i {
h.Swap(i, n)
if !down(h, i, n) {
up(h, i)
}
}
return h.Pop()
}
// Fix re-establishes the heap ordering after the element at index i has changed its value.
// Changing the value of the element at index i and then calling Fix is equivalent to,
// but less expensive than, calling Remove(h, i) followed by a Push of the new value.
// The complexity is O(log n) where n = h.Len().
func Fix(h Interface, i int) {
if !down(h, i, h.Len()) {
up(h, i)
}
}
func up(h Interface, j int) {
for {
i := (j - 1) / 2 // parent
if i == j || !h.Less(j, i) {
break
}
h.Swap(i, j)
j = i
}
}
func down(h Interface, i0, n int) bool {
i := i0
for {
j1 := 2*i + 1
if j1 >= n || j1 < 0 { // j1 < 0 after int overflow
break
}
j := j1 // left child
if j2 := j1 + 1; j2 < n && h.Less(j2, j1) {
j = j2 // = 2*i + 2 // right child
}
if !h.Less(j, i) {
break
}
h.Swap(i, j)
i = j
}
return i > i0
}
分析
heap.Interface接口需要实现:
- Push(x interface{})
- Pop() interface{}
由于嵌入了sort.Interface接口,需要另外实现其声明的方法集:
- Len() int
- Less(i, j int) bool
- Swap(i, j int)
另外,标准库中也包含了如下函数:
- Push(h Interface, x interface{})
- Pop(h Interface) interface{}
分析函数的实现:
- Push()函数调用接口中的Push()方法向堆中添加元素,再调用up()函数自底向上调整堆
- Pop()函数交换堆顶与最后一个元素,调用down()函数自顶向下调整堆,最后调用接口中的Pop()方法弹出堆顶元素
func Push(h Interface, x interface{}) {
h.Push(x)
up(h, h.Len()-1)
}
func Pop(h Interface) interface{} {
n := h.Len() - 1
h.Swap(0, n)
down(h, 0, n)
return h.Pop()
}
通过分析可知:
- 在实现heap.Interface接口时,对于Push()和Pop()两个方法的实现,只需要包括元素的加入和弹出,不包括堆的调整
- Pop()函数中,先调用了Swap()方法交换堆顶与末尾元素,因此在实现Pop()方法时,需要注意返回的应是末尾元素(原来的堆顶),在实现Push()方法时也是向末尾加入新元素
- 在使用标准库时,应该调用库中的heap.Push()和heap.Pop()函数来进行堆操作,而非自己实现的方法
再看up()和down()两个主要堆调整函数:
func up(h Interface, j int) {
for {
i := (j - 1) / 2 // parent
if i == j || !h.Less(j, i) {
break
}
h.Swap(i, j)
j = i
}
}
func down(h Interface, i0, n int) bool {
i := i0
for {
j1 := 2*i + 1
if j1 >= n || j1 < 0 { // j1 < 0 after int overflow
break
}
j := j1 // left child
if j2 := j1 + 1; j2 < n && h.Less(j2, j1) {
j = j2 // = 2*i + 2 // right child
}
if !h.Less(j, i) {
break
}
h.Swap(i, j)
i = j
}
return i > i0
}
分析可知:库中实现的默认是一个小顶堆(调用Less()方法),因此若我们需要实现一个大顶堆,只需要在实现Less()方法时进行调整即可
实例
//构造大顶堆
type MaxHeap []float64
func (h MaxHeap) Len() int {
return len(h)
}
func (h MaxHeap) Less(i, j int) bool {
return h[i] > h[j] //返回较大者
}
func (h MaxHeap) Swap(i, j int) {
h[i], h[j] = h[j], h[i]
}
func (h *MaxHeap) Push(x interface{}) {
*h = append(*h, x.(float64)) //接口类型转换
}
func (h *MaxHeap) Pop() interface{} {
old := *h
n := len(old)
x := old[n-1]
*h = old[0 : n-1]
return x
}
func halveArray(nums []int) int {
n := len(nums)
h := &MaxHeap{}
heap.Init(h)
var sum, cnt int
var minus float64
for i := 0; i < n; i++ {
//统计初始数组之和
sum += nums[i]
//加入堆
heap.Push(h, float64(nums[i]))
}
//将堆顶元素减半后重新插入
for float64(sum) > 2*minus {
top := heap.Pop(h).(float64) //接口类型转换
top /= 2.0
minus += top
heap.Push(h, top)
cnt++
}
return cnt
}
type hp struct{
sort.IntSlice
}
func(h *hp)Push(v interface{}){
h.IntSlice=append(h.IntSlice,v.(int))
}
func(h* hp)Pop()interface{}{
n:=len(h.IntSlice)-1
top:=h.IntSlice[n]
h.IntSlice=h.IntSlice[:n]
return top
}
var nums=[]int{2,3,5}
func nthUglyNumber(n int) int {
h:=&hp{sort.IntSlice{1}}
vis:=map[int]struct{}{1:{}}
top:=1
for i:=1;i<=n;i++{
top=heap.Pop(h).(int)
for j:=0;j<len(nums);j++{
mul:=nums[j]*top
if _,ok:=vis[mul];!ok{
heap.Push(h,nums[j]*top)
vis[mul]=struct{}{}
}
}
}
return top
}