2021-05-29：最常使用的K个单词II。在实时数据流中找到最常使用的k个单词，实现TopK类中的三个方法: TopK(k

package main

import "fmt"

func main() {
    a := NewTopK(2)
    a.add("lint")
    a.add("code")
    a.add("code")
    fmt.Println(a.topk())
}

type TopK struct {
    //堆
    heap     []*Node
    heapSize int
    //字，次数
    wordNodeMap map[string]*Node
    //反向表
    nodeIndexMap map[*Node]int
}

func NewTopK(k int) *TopK {
    ret := &TopK{}
    ret.heap = make([]*Node, k)

    return ret
}

func (this *TopK) add(word string) {
    if len(this.heap) == 0 {
        return
    }
    var curNode *Node
    preIndex := -1
    curNode = this.wordNodeMap[word]
    if curNode == nil {
        curNode = &Node{word, 1}
        this.wordNodeMap[word] = curNode
        this.nodeIndexMap[curNode] = -1
    } else {
        //tree set
        curNode.Times++
        preIndex = this.nodeIndexMap[curNode]
    }
    if preIndex == -1 {
        if this.heapSize == len(this.heap) {
            //treeset

        } else {
            //tree add
            this.nodeIndexMap[curNode] = this.heapSize
            this.heap[this.heapSize] = curNode
            this.HeapUp(preIndex)
        }
    } else {
        //tree add
        this.HeapDown(preIndex)
    }
}

func (this *TopK) topk() []string {
    ans := make([]string, len(this.heap))
    return ans
}

type Node struct {
    Str   string
    Times int
}

//索引上移，大根堆
func (this *TopK) HeapUp(index int) {
    for this.heap[(index-1)/2].Times < this.heap[index].Times { //父节点小于当前节点，当前节点必须上移
        this.heap[index], this.heap[(index-1)/2] = this.heap[(index-1)/2], this.heap[index]
        //加强堆
        this.nodeIndexMap[this.heap[index]], this.nodeIndexMap[this.heap[(index-1)/2]] = (index-1)/2, index
        index = (index - 1) / 2
    }
}

//索引下沉，大根堆
func (this *TopK) HeapDown(index int) {
    left := 2*index + 1
    for left <= this.heapSize-1 { //左孩子存在
        //获取大孩子
        largest := left
        if left+1 <= this.heapSize-1 && this.heap[left+1].Times > this.heap[left].Times {
            largest++
        }

        //比较
        if this.heap[index].Times < this.heap[largest].Times { //当前小于最大孩子，必须下沉
            this.heap[index], this.heap[largest] = this.heap[largest], this.heap[index]
            //加强堆
            this.nodeIndexMap[this.heap[index]], this.nodeIndexMap[this.heap[largest]] = largest, index
        } else {
            break
        }

        //下一次遍历
        index = largest
        left = 2*index + 1
    }
}

func (this *TopK) Push(node *Node) {
    this.heap[this.heapSize] = node
    //加强堆
    this.nodeIndexMap[node] = this.heapSize
    this.heapSize++
    //索引上移
    this.HeapUp(this.heapSize)
}

func (this *TopK) Pop() *Node {
    ans := this.heap[0]
    this.heap[0], this.heap[this.heapSize-1] = this.heap[this.heapSize-1], this.heap[0]
    //加强堆
    this.nodeIndexMap[this.heap[0]] = 0
    this.nodeIndexMap[this.heap[this.heapSize-1]] = -1
    this.heapSize--
    //索引下沉
    this.HeapDown(0)
    return ans
}