最小堆 以及 优先级队列的 Golang 实现

[JemmyH]

堆，是计算机科学中的一种特别的完全二叉树。若父节点的值恒小于等于子节点的值，此堆称为最小堆（min heap）；反之，若母节点的值恒大于等于子节点的值，此堆称为最大堆（max heap）。在堆中最顶端的那一个节点，称作 根节点（root node），根节点本身没有 父节点（parent node）。堆通常是一个可以被看做一棵树的数组对象。在队列中，调度程序反复提取队列中第一个作业并运行，因为实际情况中某些时间较短的任务将等待很长时间才能结束，或者某些不短小，但具有重要性的作业，同样应当具有优先权。堆即为解决此类问题设计的一种数据结构。

优先级队列 是计算机科学中的一类抽象数据类型。优先队列中的每个元素都有各自的优先级，优先级最高的元素最先得到服务；优先级相同的元素按照其在优先队列中的顺序得到服务。优先队列往往用堆来实现。

[JemmyH]

Golang实现一：根据原理简单实现：

package minheap

import (
	"container/heap"
	"fmt"
	"math"

	"github.com/pkg/errors"
)

/*
* @CreateTime: 2021/7/6 21:57
* @Author: hujiaming
* @Description: Golang实现最小堆
 */

var ErrMinHeapEmpty = errors.New("minHeap is empty")

const HeapHeadTag int64 = math.MinInt64

type MinHeap struct {
	elements []int64
}

// NewMinHeap 创建一个最小堆实例
func NewMinHeap() *MinHeap {
	return &MinHeap{elements: []int64{HeapHeadTag}}
}

/*
Add 将一个元素添加到最小堆中，并且添加后要使其满足最小堆的特性
首先将该元素插入到数组最后，然后对这最后一个元素进行 “上浮” 操作：
该元素与父元素进行大小比较，如果小于父元素，则和父元素交换位置，如此循环，直到 到达堆顶 或 子元素小于父元素。
*/
func (mh *MinHeap) Add(v int64) {
	// 1. 先将元素插在数组最后面
	mh.elements = append(mh.elements, v)
	// 2. 将最后一个元素上浮，使其符合最小堆的性质。其实是为 v 找位置
	i := len(mh.elements) - 1
	for ; mh.elements[i/2] > v; i /= 2 {
		mh.elements[i] = mh.elements[i/2]
	}
	mh.elements[i] = v
}

/*
PopMin 弹出堆中最小的元素
对最小堆而言，移除元素，只能移除堆顶(最小值)的元素。
首先，移除堆顶元素，然后将最后一个元素放在堆顶，之后对这第一个元素进行 “下沉” 操作：
将此元素与两个子节点元素比较，如果当前结点大于两个子节点，则与较小的子节点交换位置，如此循环，直到 到达叶子结点 或 小于较小子节点。
*/
func (mh *MinHeap) PopMin() (int64, error) {
	if mh.IsEmpty() {
		return 0, ErrMinHeapEmpty
	}
	res := mh.elements[1]
	last := mh.elements[len(mh.elements)-1]
	// idx 表示最后一个元素应该在的位置
	var idx int
	for idx = 1; idx*2 < len(mh.elements); {
		// 找出子节点中较小的元素的 index
		minChildIdx := idx * 2
		if minChildIdx < len(mh.elements)-1 && mh.elements[minChildIdx+1] < mh.elements[minChildIdx] {
			minChildIdx++
		}
		// 当前结点 大于 较小子节点，和这个较小子节点交换位置，继续循环
		if last > mh.elements[minChildIdx] {
			mh.elements[idx] = mh.elements[minChildIdx]
			idx = minChildIdx
			continue
		}
		break
	}
	mh.elements[idx] = last
	mh.elements = mh.elements[:len(mh.elements)-1]

	return res, nil
}

// PeekHead 只返回堆顶元素(最小值)，不进行下沉操作
func (mh *MinHeap) PeekHead() (int64, error) {
	if mh.IsEmpty() {
		return 0, ErrMinHeapEmpty
	}
	return mh.elements[1], nil
}

// IsEmpty 最小堆是否是空的
func (mh *MinHeap) IsEmpty() bool {
	if len(mh.elements) == 0 || (len(mh.elements) == 1 && mh.elements[0] == HeapHeadTag) {
		return true
	}
	return false
}

// Length 返回最小堆中的元素个数
func (mh *MinHeap) Length() int {
	return len(mh.elements) - 1
}

// Print 打印代表最小堆的数组
func (mh *MinHeap) Print() {
	fmt.Println(mh.elements[1:])
}

Test 如下：

func TestMinHeap(t *testing.T) {
	mh := NewMinHeap()
	mh.Add(4)
	mh.Add(2)
	mh.Add(7)
	mh.Add(9)
	mh.Add(1)
	mh.Add(5)
	mh.Add(10)
	mh.Add(3)
	mh.Add(2)
	mh.Print()
	for !mh.IsEmpty() {
		fmt.Println(mh.PopMin())
	}
	assert.Equal(t, mh.Length(), 0)
}

// 输出
/*

[1 2 5 2 4 7 10 9 3]
1 <nil>
2 <nil>
2 <nil>
3 <nil>
4 <nil>
5 <nil>
7 <nil>
9 <nil>
10 <nil>

*/

[JemmyH]

Golang 实现二：实现标准库 heap.Interface 接口

先看下标准库中的 Interface，位置在 container/heap/heap.go：

// 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.
}

// An implementation of Interface can be sorted by the routines in this package.
// The methods refer to elements of the underlying collection by integer index.
type Interface interface {
	// Len is the number of elements in the collection.
	Len() int

	// Less reports whether the element with index i
	// must sort before the element with index j.
	//
	// If both Less(i, j) and Less(j, i) are false,
	// then the elements at index i and j are considered equal.
	// Sort may place equal elements in any order in the final result,
	// while Stable preserves the original input order of equal elements.
	//
	// Less must describe a transitive ordering:
	//  - if both Less(i, j) and Less(j, k) are true, then Less(i, k) must be true as well.
	//  - if both Less(i, j) and Less(j, k) are false, then Less(i, k) must be false as well.
	//
	// Note that floating-point comparison (the < operator on float32 or float64 values)
	// is not a transitive ordering when not-a-number (NaN) values are involved.
	// See Float64Slice.Less for a correct implementation for floating-point values.
	Less(i, j int) bool

	// Swap swaps the elements with indexes i and j.
	Swap(i, j int)
}

我们以此为基础，实现一个 优先级队列:

package priorityqueen


type Item struct {
	value    int64 // 实际值
	priority int64 // 优先级
	index    int   // 当前 item 在数组中的 index
}

// PriorityQueen 表示优先级队列
type PriorityQueen []*Item

func (mh2 PriorityQueen) Len() int {
	return len(mh2)
}

func (mh2 PriorityQueen) Less(i, j int) bool {
	return mh2[i].priority < mh2[j].priority
}

func (mh2 PriorityQueen) Swap(i, j int) {
	mh2[i], mh2[j] = mh2[j], mh2[i]
	mh2[i].index = i
	mh2[j].index = j
}

// Push 将 x 添加到数组最后
func (mh2 *PriorityQueen) Push(x interface{}) {
	l := len(*mh2)
	c := cap(*mh2)
	if l+1 > c {
		cmh2 := make([]*Item, l, c/2)
		copy(*mh2, cmh2)
		*mh2 = cmh2
	}
	*mh2 = (*mh2)[:l+1]
	item := (x).(*Item)
	item.index = l
	(*mh2)[l] = item
}

// Pop 返回数组最后一个元素
func (mh2 *PriorityQueen) Pop() interface{} {
	l := len(*mh2)
	c := cap(*mh2)
	if l < c/2 && c > 25 {
		cmh2 := make([]*Item, l, c/2)
		copy(cmh2, *mh2)
		*mh2 = cmh2
	}
	item := (*mh2)[l-1]
	item.index = -1 // for safety
	*mh2 = (*mh2)[:l-1]
	return item
}

// PopHead 弹出堆顶元素
func (mh2 *PriorityQueen) PopHead() *Item {
	if mh2.Len() == 0 {
		return nil
	}
	item := (*mh2)[0]

	heap.Remove(mh2, 0)

	return item
}

// PopWithPriority 弹出优先级小于 maxP 的堆顶元素，如果没有，返回 nil 和 当前堆顶和maxP的距离
func (mh2 *PriorityQueen) PopWithPriority(maxP int64) (*Item, int64) {
	if mh2.Len() == 0 {
		return nil, 0
	}
	item := (*mh2)[0]
	if item.priority > maxP {
		return nil, item.priority - maxP
	}

	heap.Remove(mh2, 0)

	return item, 0
}

// PeekHead 显示堆顶元素
func (mh2 *PriorityQueen) PeekHead() *Item {
	if mh2.Len() == 0 {
		return nil
	}
	heap.Init(mh2)
	item := (*mh2)[0]
	return item
}

测试一下：

func TestPriorityQueen(t *testing.T) {
	items := make([]*Item, 0)
	rand.Seed(time.Now().UnixNano())

	for i := 0; i < 10; i++ {
		v := rand.Int63n(100)
		items = append(items, &Item{
			value:    v,
			priority: v,
			index:    i,
		})
	}
	q := PriorityQueen(items)
	heap.Init(&q)

	fmt.Println(q.PeekHead())

	maxP := int64(50)
	for _, i := range q {
		if i.priority < maxP {
			fmt.Println(fmt.Sprintf("p: %d, v: %d", i.priority, i.value))
		}
	}

	fmt.Println("====")
	for i := 0; i < 10; i++ {
		item, _ := q.PopWithPriority(maxP)
		if item != nil {
			fmt.Println(item)
		}
	}

	fmt.Println("====")
	for {
		item := q.PopHead()
		if item == nil {
			break
		}
		fmt.Println(item)
	}
}

// 输出
/*

&{5 5 0}
p: 5, v: 5
p: 11, v: 11
p: 6, v: 6
p: 33, v: 33
====
&{5 5 -1}
&{6 6 -1}
&{11 11 -1}
&{33 33 -1}
&{50 50 -1}
====
&{52 52 -1}
&{73 73 -1}
&{85 85 -1}
&{97 97 -1}
&{99 99 -1}

*/

[JemmyH]

Golang 标准库 heap.Interface 源码解析

整个包的实现非常简洁，加上注释以及空行，整个文件才只有120 行：

// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package heap provides heap operations for any type that implements
// heap.Interface. A heap is a tree with the property that each node is the
// minimum-valued node in its subtree.
//
// The minimum element in the tree is the root, at index 0.
//
// A heap is a common way to implement a priority queue. To build a priority
// queue, implement the Heap interface with the (negative) priority as the
// ordering for the Less method, so Push adds items while Pop removes the
// highest-priority item from the queue. The Examples include such an
// implementation; the file example_pq_test.go has the complete source.
//
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()
        // (n/2 - 1) 处的结点是最后一棵子树(没有孩子结点)的根节点
	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
}

我们关注其中几个核心实现：

• down(h Interface, idx, heapLen int) 下沉操作：
  首先，移除堆顶元素，然后将最后一个元素放在堆顶，之后对这第一个元素进行 “下沉” 操作：
  将此元素与两个子节点元素比较，如果当前结点大于两个子节点，则与较小的子节点交换位置，如此循环，直到 到达叶子结点 或 小于较小子节点。
  为什么元素 i 比它的两个子节点都小，就可以跳出循环，不再继续下去呢？这是由于，在 Init 函数中，**第一个开始 down 的元素是第 n/2 - 1 个，可以保证总是从最后一棵子树开始 down **，因此可以保证 Init->down 时，如果元素 i 比它的两个子节点都小，那么该元素对应的子树，就是最小堆。
  

• up(h Interface, curIdx int) 上浮操作：
  主要用在 Push 中，当我们向最小堆插入一个元素时，现将其插入到数组最后，之后进行上浮操作，此时的 curIdx 就是数组最后一个元素的 index，即 h.Len() - 1。当前元素与其父元素进行比较，如果当前元素小于父元素，则与父元素交换位置，如此往复，直到堆顶或者当前元素大于父元素。