LC347 - 前K个高频元素

package Heap

import (
    "container/heap"
)

// Comparable is an interface.
type Comparable[T interface{}] interface {
    MyLess(other T) bool
}

// MyHeap is a generic type.
// The elements could be user-defined data types.
// You should also implement the interface 'Comparable'.
type MyHeap[T Comparable[T]] []T

func (h MyHeap[T]) Len() int {
    return len(h)
}

func (h MyHeap[T]) Less(i, j int) bool {
    return h[i].MyLess(h[j])
}

func (h MyHeap[T]) Swap(i, j int) {
    h[i], h[j] = h[j], h[i]
}

func (h *MyHeap[T]) Push(x interface{}) {
    *h = append(*h, x.(T))
}

func (h *MyHeap[T]) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[0 : n-1]
    return x
}

func (h *MyHeap[T]) Top() interface{} {
    if h.Len() == 0 {
        return nil
    } else {
        return (*h)[0]
    }
}

type Frequency struct {
    val, freq int
}

func (cur *Frequency) MyLess(other *Frequency) bool {
    return cur.freq < other.freq
}

func topKFrequent(nums []int, k int) []int {
    hashTable := make(map[int]int)
    for _, num := range nums {
        hashTable[num]++
    }
    minHeap := new(MyHeap[*Frequency])
    heap.Init(minHeap)
    for val, freq := range hashTable {
        if minHeap.Len() < k {
            heap.Push(minHeap, &Frequency{val, freq})
        } else if freq > minHeap.Top().(*Frequency).freq {
            heap.Push(minHeap, &Frequency{val, freq})
            heap.Pop(minHeap)
        }
    }
    result := make([]int, 0)
    for minHeap.Len() > 0 {
        result = append(result, heap.Pop(minHeap).(*Frequency).val)
    }
    return result
}

本题的逻辑和上题LC215几乎没有区别，需要注意的是如何利用泛型接口实现自定义数据类型的堆结构。