Skip to content

Instantly share code, notes, and snippets.

@alldroll
Created July 2, 2020 12:29
Show Gist options
  • Save alldroll/e43850442d5c80bbe3b24143a3ee31d5 to your computer and use it in GitHub Desktop.
Save alldroll/e43850442d5c80bbe3b24143a3ee31d5 to your computer and use it in GitHub Desktop.
import "container/heap"
type Item struct {
word string
count int
}
func (i Item) Less(o Item) bool {
if i.count == o.count {
return i.word > o.word
}
return i.count < o.count
}
type Heap []Item
func (h Heap) Len() int { return len(h) }
func (h Heap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h Heap) Less(i, j int) bool {
return h[i].Less(h[j])
}
func (h *Heap) Push(x interface{}) {
*h = append(*h, x.(Item))
}
func (h *Heap) Pop() interface{} {
n := h.Len()
old := *h
x := old[n - 1]
*h = old[:n - 1]
return x
}
func topKFrequent(words []string, k int) (result []string) {
counts := make(map[string]int)
for _, word := range words {
counts[word]++
}
h := Heap{}
for word, count := range counts {
item := Item{word, count}
if h.Len() < k {
heap.Push(&h, item)
} else if h[0].Less(item) {
h[0] = item
heap.Fix(&h, 0)
}
}
for h.Len() > 0 {
item := heap.Pop(&h).(Item)
result = append(result, item.word)
}
reverse(result)
return result
}
func reverse(arr []string) {
for i, j := 0, len(arr) - 1; i < j; i, j = i + 1, j - 1 {
arr[i], arr[j] = arr[j], arr[i]
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment