go语言实现大根堆、小根堆及堆排序

go 评论

二叉堆是一种特殊的堆，它满足两个性质：结构性和堆序性

  • 结构性：二叉堆是一棵完全二叉树，完全二叉树可以用一个数组表示，不需要指针，所以效率更高。当用数组表示时，数组中任一位置i上的元素，其左子树在位置2i上，右子树在位置2i+1上，其父节点在位置i/2上。
  • 堆序性质：堆的最小值或最大值在根节点上，所以可以快速找到最大值或最小值。

最大堆和最小堆是二叉堆的两种形式：

  • 最大堆：根节点的键值是所有堆节点键值中最大者的堆。
  • 最小堆：根节点的键值是所有堆节点键值中最小者的堆。

最小堆实现

插入和删除

当向最小堆插入元素时：

  • 将元素插入末尾
  • 判断该元素是否需要上移(与父节点比较，如果比父节点小则上移)
  • 重复上述步骤，直到满足最小堆特性

当向最小堆删除元素时：

  • 删除堆顶元素
  • 判断目前的堆顶元素是否需要下调(与子节点比较，和其中较小的节点交换位置)
  • 重复上述步骤，直到满足最小堆特性

具体实现

下面以求数据流中第k大的元素为问题实现一个最小堆，如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
type minHeap struct {
    k    int   // 容量
    heap []int // heap数组
}

func createMinHeap(k int, nums []int) *minHeap {
    heap := &minHeap{k: k, heap: []int{}}
    for _, n := range nums {
        heap.add(n)
    }
    return heap
}

func (m *minHeap) add(num int) {
    if len(m.heap) < m.k {
        m.heap = append(m.heap, num)
        m.up(len(m.heap) - 1)
    } else if num > m.heap[0] {
        m.heap[0] = num
        m.down(0)
    }
}

// 元素上浮
func (m *minHeap) up(i int) {
    for i > 0 {
        parent := (i - 1) >> 1 // 找到父节点在heap数组中的位置
        // 如果比父节点元素小，则交换位置并更新索引
        if m.heap[parent] > m.heap[i] {
            m.heap[parent], m.heap[i] = m.heap[i], m.heap[parent]
            i = parent
        } else {
            break // 当前节点比父节点小，满足最小堆性质，退出
        }
    }
}

// 元素下沉(包括切片中第一个元素，索引为0)
func (m *minHeap) down(i int) {
    for 2*i+1 < len(m.heap) { // 左子节点越界，则退出循环
        child := 2*i + 1 // 左子节点在heap切片中的位置
        if child+1 < len(m.heap) && m.heap[child+1] < m.heap[child] {
            child++ // 如果右子节点没有越界，且值比左子节点更小，则选择下沉右子节点
        }
        // 将当前元素与子节点最大元素对比，然后交换并更新索引
        if m.heap[i] > m.heap[child] {
            m.heap[child], m.heap[i] = m.heap[i], m.heap[child]
            i = child
        } else {
            break // 子节点都比自己大，满足最小堆属性，退出
        }
    }
}

应用

如果要求输出数据流中的第k大元素，正好可以使用最小堆实现：

1
2
3
4
5
6
7
8
9
10
11
12
type KthLargest struct {
    heap *minHeap
}

func Constructor(k int, nums []int) KthLargest {
    return KthLargest{heap: createMinHeap(k, nums)}
}

func (k *KthLargest) Add(val int) int {
    k.heap.add(val)
    return k.heap.heap[0]
}

使用heap包实现

heap源码中定义了一个Interface接口，该接口一共包含5个方法，定义一个实现了该接口的结构体就实现了一个二叉堆。
container/heap/heap.go

1
2
3
4
5
6
7
8
// 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.
}

sort/sort.go

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
// 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)
}

定义一个最大堆，实现上述接口如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
type MaxHeap []int

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.(int))
}

// Pop 弹出堆顶元素
func (h *MaxHeap) Pop() interface{} {
    res := (*h)[len(*h)-1]
    *h = (*h)[:len(*h)-1]
    return res
}

func main() {
    h := &MaxHeap{3, 1, 2, 5}
    heap.Init(h)
    heap.Push(h, 8)
    for h.Len() > 0 {
        fmt.Printf("%d ", heap.Pop(h))
    }
}

堆排序

堆排序是一种选择排序，它的最坏、最好、平均时间复杂度均为O(nlogn)，它也是不稳定排序。

  • 排序的过程主要由构建初始堆，交换堆顶元素和末尾元素并重建堆两部分组成
  • 升序使用最大堆，每次和末尾元素交换，然后重新构建最大堆，整体数组减一；反之降序使用最小堆
  • 堆构建从第一个非叶子节点开始，也就是len/2 - 1所在位置的元素
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
func maxHeap(nums []int, length int) []int {
    if length <= 1 {
        return nums
    }
    parent := length/2 + 1 // 第一个非叶子节点
    for i := parent; i >= 0; i-- {
        // 比较三个节点的大小并将较大的节点上浮
        max := i
        leftChild := 2*i + 1
        rightChild := 2*i + 2
        if leftChild <= length-1 && nums[leftChild] > nums[max] {
            max = leftChild
        }
        if rightChild <= length-1 && nums[rightChild] > nums[max] {
            max = rightChild
        }
        if max != i {
            nums[i], nums[max] = nums[max], nums[i]
        }
    }
    return nums
}

func sortHeap(nums []int) []int {
    length := len(nums)
    for i := 0; i < length; i++ {
        lastLength := length - i // 剔除已经排完序的元素
        nums = maxHeap(nums, lastLength)  // 重新构建最大堆
        nums[0], nums[lastLength-1] = nums[lastLength-1], nums[0]
    }
    return nums
}

func main() {
    nums := []int{8, 5, 11, 2, 7, 9}
    fmt.Println(sortHeap(nums))
}

HAHADeveloper & Designer

个人简介
Powered By Valine
v1.5.2