347. 前 K 个高频元素【最小堆】

题目

代码

import (
    "bytes"
    "fmt"
)

// 输入: nums = [1,1,1,2,2,3], k = 2
// 输出: [1,2]
func topKFrequent(nums []int, k int) []int {
    // 统计各个元素出现的次数
    eleFreqMap := make(map[int]*EleFreq)
    for _, num := range nums {
        if _, ok := eleFreqMap[num]; !ok {
            eleFreqMap[num] = &EleFreq{num, 0}
        }
        eleFreqMap[num][Freq]++
    }
    //fmt.Printf("eleFreqMap is %v\n\n", eleFreqMap)

    var kCount int
    mh := NewMinHeap()
    for i := range eleFreqMap {
        // 建堆
        if kCount < k {
            mh.SiftUp(eleFreqMap[i])
            kCount ++
            continue
        } else {
            // 小于没有机会
            minVal := mh.GetMinFreq()
            if eleFreqMap[i][Freq] > minVal {
                mh.Replace(eleFreqMap[i])       // 替换
            }
        }
    }
    //fmt.Printf("mh is %v\n\n", mh)

    ans := make([]int, 0, k)
    // 出堆
    for i:=0; i<k; i++ {
        ans = append(ans, mh.SiftDown())
    }

    return ans
}

// 次数统计
const (
    Ele  = 0
    Freq = 1
)

type EleFreq [2]int // [0]元素 [1]出现次数
func(ef EleFreq) String() string {
    return fmt.Sprintf("<ele: %v, freq: %v>", ef[Ele], ef[Freq])
}

// 最小堆
type MinHeap struct {
    Size int
    EFs  []*EleFreq
}

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

func(mh MinHeap) String() string {
    var strBuf bytes.Buffer

    strBuf.WriteString("MinHeap:<\n")
    strBuf.WriteString(fmt.Sprintf("Size: %v\n", mh.Size))
    strBuf.WriteString("EFs: {")
    for k := range mh.EFs {
        strBuf.WriteString(mh.EFs[k].String())
    }
    strBuf.WriteString("}\n>\n")

    return strBuf.String()
}

// 获得父节点下标
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 (mh *MinHeap) SiftUp(ef *EleFreq) {
    // 插在数组的尾部，也就是Size处
    childIdx := mh.Size
    parentIdx := GetParentIdx(childIdx)
    mh.EFs = append(mh.EFs, ef)

    // 最小堆，根节点小于子节点
    for mh.EFs[parentIdx][Freq] > mh.EFs[childIdx][Freq] {
        mh.EFs[parentIdx], mh.EFs[childIdx] = mh.EFs[childIdx], mh.EFs[parentIdx]
        childIdx = parentIdx
        parentIdx = GetParentIdx(parentIdx) // 上浮
    }

    mh.Size++
}

// 删除时下沉
// 此函数不做删除操作，只做下沉操作，参数是根节点的索引
func (mh *MinHeap) siftDown(rootIdx int) {

    var minChildIdx int // 左右子节点中的较小者

    for {
        leftChildIdx := GetLeftChildIdx(rootIdx)
        rightChildIdx := GetRightChildIdx(rootIdx)

        switch { // 只会进入1个case，不会穿透
        case leftChildIdx > mh.Size-1: // 左子节点超出范围，（肯定没有右子节点）
            return // 终止循环
        case rightChildIdx > mh.Size-1: // 左子节点未超出范围，但是右子几点超出访问
            if mh.EFs[rootIdx][Freq] > mh.EFs[leftChildIdx][Freq] {
                mh.EFs[rootIdx], mh.EFs[leftChildIdx] = mh.EFs[leftChildIdx], mh.EFs[rootIdx]
            }
            return // 终止循环

            // 寻找较小的子节点
        case mh.EFs[leftChildIdx][Freq] <= mh.EFs[rightChildIdx][Freq]:
            minChildIdx = leftChildIdx

        case mh.EFs[leftChildIdx][Freq] > mh.EFs[rightChildIdx][Freq]:
            minChildIdx = rightChildIdx
        }

        // 交换
        if mh.EFs[rootIdx][Freq] > mh.EFs[minChildIdx][Freq] {
            mh.EFs[rootIdx], mh.EFs[minChildIdx] = mh.EFs[minChildIdx], mh.EFs[rootIdx]
            rootIdx = minChildIdx // 对较小子节点进行同样的操作
        } else {
            return  // 终止循环
        }
    }
}

// 删除堆顶元素，并将数组尾部的元素放在堆顶，然后执行下沉操作
// 为了最后的排序输出
func (mh *MinHeap) SiftDown() int {
    // 删除操作，Size--
    mh.Size --
    // 此时的队尾元素所在的下标是Size
    minVal := mh.GetMinEle()
    mh.EFs[0], mh.EFs[mh.Size] = mh.EFs[mh.Size], mh.EFs[0]

    mh.siftDown(0)  // 从根节点开始执行下沉操作
    return minVal
}

// 将堆顶元素替换为指定的元素，执行下沉操作
func (mh *MinHeap) Replace(ef *EleFreq) {
    mh.EFs[0] = ef
    mh.siftDown(0)  // 从根节点开始执行下沉操作
}

func (mh *MinHeap) GetMinFreq() int {
    return mh.EFs[0][Freq]
}

func (mh *MinHeap) GetMinEle() int {
    return mh.EFs[0][Ele]
}