【堆】【最小堆】

堆

// 堆实际上是一个完全二叉树
// 如果用0是堆的根节点，那么：
// 最子节点是20+1，右子节点是20+2
//
// 如果i是节点编号，那么其父节点是：
// i是奇数 (i-1)/2
// i是偶数(i-2)/2
//
// 因此，可以使用数组来实现堆！
//
// 大顶堆，就是 根节点最大，并且子树也是大顶堆
// 小顶堆，就是 根节点最小，并且子树也是小顶堆
//
// 所以，计算前k个最大数，使用小顶堆，因为堆中存放的是“当前”最大的k个数，后面的数，如果小于堆中最小数，就没有机会进入堆，如果大于最小数，则需要更新堆！
// 同理，计算前k个最小数，使用大顶堆。

// 从数据的存储结构看，最大堆/最小堆是一个数组
// 从数据的逻辑结构看，最大堆/最小堆是一个完全二叉树
//
// 最小堆的实现
// 插入：
// 插入元素时，插入到数组中的最后一个元素的后面，然后与该节点的父节点比较大小。
// 如果插入的元素小于父节点元素，那么与父节点交换位置。然后插入元素交换到父节点位置时，又与该节点的父节点比较，直到大于父节点元素或者到达堆顶。
// 该过程叫做上浮，即插入时上浮。
//
// 删除：
// 移除元素时，只能从堆顶移除元素，再取最后一个元素放到堆顶中。
// 然后堆顶节点与子节点比较时，先取子节点中的较小者，如果堆顶节点大于较小子节点，那么交换位置。此时堆顶节点元素交换到较小子节点上。然后再与其较小子节点比较，直到小于较小子节点或者到达叶子节点为止。
// 该过程叫做下沉，即移除元素时下沉。
//
/*
0
1 2
3 4 5
*/

代码实现

package main

import (
    "errors"
    "fmt"
)

// 堆实际上是一个完全二叉树
// 如果用0是堆的根节点，那么：
// 最子节点是2*0+1，右子节点是2*0+2
//
// 如果i是节点编号，那么其父节点是：
// i是奇数 (i-1)/2
// i是偶数(i-2)/2
//
// 因此，可以使用数组来实现堆！
//
// 大顶堆，就是 根节点最大，并且子树也是大顶堆
// 小顶堆，就是 根节点最小，并且子树也是小顶堆
//
// 所以，计算前k个最大数，使用小顶堆，因为堆中存放的是“当前”最大的k个数，后面的数，如果小于堆中最小数，就没有机会进入堆，如果大于最小数，则需要更新堆！
//       同理，计算前k个最小数，使用大顶堆。

// 从数据的存储结构看，最大堆/最小堆是一个数组
// 从数据的逻辑结构看，最大堆/最小堆是一个完全二叉树
//
// 最小堆的实现
// 插入：
// 插入元素时，插入到数组中的最后一个元素的后面，然后与该节点的父节点比较大小。
// 如果插入的元素小于父节点元素，那么与父节点交换位置。然后插入元素交换到父节点位置时，又与该节点的父节点比较，直到大于父节点元素或者到达堆顶。
// 该过程叫做上浮，即插入时上浮。
//
// 删除：
// 移除元素时，只能从堆顶移除元素，再取最后一个元素放到堆顶中。
// 然后堆顶节点与子节点比较时，先取子节点中的较小者，如果堆顶节点大于较小子节点，那么交换位置。此时堆顶节点元素交换到较小子节点上。然后再与其较小子节点比较，直到小于较小子节点或者到达叶子节点为止。
// 该过程叫做下沉，即移除元素时下沉。
//
/*
        0
     1      2
   3   4  5
*/
// 最小堆的实现
type MinHeap struct {
    size int
    nums []int
}

func NewMinHeap() *MinHeap {
    return &MinHeap{}
}

// 获得父节点下标
func GetParentIdx(i int) int {
    if i == 0 {
        return 0 // 根节点
    }
    return (i - 1) / 2
}

// 获得左子节点下标
func GetLeftChildIdx(i int) int {
    return 2*i + 1
}

// 获得右子节点下标
func GetRightChildIdx(i int) int {
    return 2*i + 2
}

func (h *MinHeap) isEmpty() bool {
    return h.size == 0
}

// 获得最小数，既堆顶元素
func (h *MinHeap) getMinValue() (int, error) {
    if h.isEmpty() {
        return 0, errors.New("getMinValue failed, minHeap is empty")
    }

    return h.nums[0], nil
}

// 插入时上浮
// 参数num，是插入的数字
func (h *MinHeap) SiftUp(num int) {
    fmt.Printf("SiftUp, nums is %v\n", num)
    h.nums = append(h.nums, num)
    childIndex := h.size // 开始的时候，插入到数组最后
    parentIndex := GetParentIdx(childIndex)

    // 最小堆，父节点要小于子节点
    for h.nums[parentIndex] > h.nums[childIndex] {
        h.nums[parentIndex], h.nums[childIndex] = h.nums[childIndex], h.nums[parentIndex]
        childIndex = parentIndex
        parentIndex = GetParentIdx(parentIndex) // 上浮
    }

    h.size++
}

// 然后堆顶节点与子节点比较时，先取子节点中的较小者，如果堆顶节点大于较小子节点，那么交换位置。
// 从rootIdx节点开始进行下沉操作，不删除节点
func (h *MinHeap) siftDown(rootIdx int) {
    //fmt.Printf("SiftDown, rootIdx is %v\n", rootIdx)
    var minChildIdx int
    for {
        leftChildIdx := GetLeftChildIdx(rootIdx)
        rightChildIdx := GetRightChildIdx(rootIdx)

        //fmt.Printf("rootIdx is %v, siftDown, leftChildIdx is %v, rightChildIdx is %v, minChildIdx is %v\n", rootIdx, leftChildIdx, rightChildIdx, minChildIdx)

        // go case不会穿透
        switch {
        case leftChildIdx > h.size-1: // 左子节点下标，已经超出堆的大小范围
            return
        case rightChildIdx > h.size - 1: // 左子节点下标，已经超出堆的大小范围(此时，左子节点仍然在范围内)
            if h.nums[rootIdx] > h.nums[leftChildIdx] {
                h.nums[rootIdx], h.nums[leftChildIdx] = h.nums[leftChildIdx], h.nums[rootIdx]  // 跟左子节点交换
            }
            return

            // 下面两个case寻找子节点中比较小的
        case h.nums[leftChildIdx] <= h.nums[rightChildIdx]:
            minChildIdx = leftChildIdx
        case h.nums[leftChildIdx] > h.nums[rightChildIdx]:
            minChildIdx = rightChildIdx
        }

        // 如果堆顶节点大于较小子节点，那么交换位置
        if h.nums[rootIdx] > h.nums[minChildIdx] {
            h.nums[rootIdx], h.nums[minChildIdx] = h.nums[minChildIdx], h.nums[rootIdx]    // 与较小者进行交换
            //fmt.Printf("rootIdx is %v, minChildIdx is %v\n", rootIdx, minChildIdx)
        }

        rootIdx = minChildIdx                                                          // 对较小的子节点进行相同的操作
    }
}

// 删除堆顶操作，将最后一个元素放到堆顶，进行下沉操作
func (h *MinHeap) SiftDown() (int, error) {
    minVal, err := h.getMinValue()
    if err != nil {
        return minVal, err
    }

    h.size --
    h.nums[0], h.nums[h.size] = h.nums[h.size], h.nums[0]

    h.siftDown(0)   // 从堆顶开始进行下沉操作
    return minVal, err
}

// 用指定的num，替换堆顶元素
// 然后对根节点进行下沉操作
func(h *MinHeap) Replace(num int) (int, error) {
    minVal, err := h.getMinValue()
    if err != nil {
        return minVal, err
    }

    h.nums[0] = num
    h.siftDown(0)
    return minVal, err
}

// 求前k个最大元素
func MaxKNums(nums []int, k int) (rsp []int) {
    // 先用前k个元素，构造堆
    heap := NewMinHeap()
    for i :=0; i<k; i++ {
        heap.SiftUp(nums[i])   // 建堆就是，将元素放入数组末尾，然后进行上浮操作
    }

    minVal, _ := heap.getMinValue()
    for _, v := range nums[k:] {
        // 如果比堆顶元素还小，你就肯定不够资格成为前k大元素
        if v > minVal {
            heap.Replace(v)                // 入堆
            minVal, _ = heap.getMinValue() // 重新获取堆顶值
        }
    }

    // 将堆中元素排好序，输出，每次删除堆顶元素（下沉操作）
    for i := 0; i<k; i++ {
        v, _ := heap.SiftDown()
        rsp = append(rsp,  v)
    }

    return
}