Dev Skills

Frequency count and top K

Count with a map, rank with a sort or a heap, and never let map order leak into your answer.

Topic 26 of 98 · 5 min to read · 55 min to build
After this topic you can
  • Write a top-k word counter with a map and slices.SortFunc.
  • Explain why a tie-break is needed for deterministic output.
  • Predict the output order when two words tie on count.

The problem

Given some text, return the k most frequent words with their counts. Ties are broken alphabetically.

Brute force first

For each word, count it by scanning the whole text again. Then repeat k times: pick the word with the highest count, print it, remove it.

Counting that way is O(n²), and picking is O(k·m). The map does the counting in one pass.

The approach

Count with a map[string]int. Maps have no order, so move the pairs into a slice and sort by count descending, then word ascending: O(n + m log m) for m distinct words.

This is the answer to give first. Most interviewers are happy with it; a faster heap version is at the end of the lesson, as an optional extra.

top_k.go
package main
import (
"cmp"
"fmt"
"slices"
"strings"
)
type wordCount struct {
word string
count int
}
// topK: count with a map, move to a slice, sort. O(n + m log m) for m distinct words.
func topK(text string, k int) []wordCount {
counts := map[string]int{}
for _, w := range strings.Fields(strings.ToLower(text)) {
counts[w]++
}
list := make([]wordCount, 0, len(counts))
for w, c := range counts {
list = append(list, wordCount{w, c})
}
// Count descending, then word ascending, so the output is the same every run.
// Without the tie-break, map iteration order leaks into the answer.
slices.SortFunc(list, func(a, b wordCount) int {
if c := cmp.Compare(b.count, a.count); c != 0 {
return c
}
return cmp.Compare(a.word, b.word)
})
return list[:min(k, len(list))]
}
func main() {
fmt.Println(topK("the cat and the dog and the bird", 2)) // [{the 3} {and 2}]
fmt.Println(topK("Go go GO rust", 5)) // [{go 3} {rust 1}]
}

Dry run

"the cat and the dog and the bird", k = 2:

1. Count: the 3, and 2, cat 1, dog 1, bird 1 (in random map order).

2. Sort by count desc, then word asc: the 3, and 2, bird 1, cat 1, dog 1.

3. Take list[:min(2, 5)]: [{the 3} {and 2}].

Go-specific gotchas

  • Map iteration order is randomised on purpose. Without a tie-break in the comparator, two runs can print different answers and the judge marks you wrong.
  • `slices.SortFunc` wants an `int` comparator (negative, zero, positive). cmp.Compare writes it for you. The older sort.Slice takes a less bool and still works.
  • `min` and `max` are built-ins since Go 1.21. list[:min(k, len(list))] stops a panic when k is bigger than the list.

Faster for big n (optional)

Skip this on a first read. If k is small (say 10) and there are hundreds of thousands of distinct words, sorting all of them is wasted work. Keep only the best k in a min-heap from container/heap: push each word, and pop the weakest when the heap grows past k. That is O(n + m log k).

`any` is just another name for interface{}: a box that can hold a value of any type. container/heap was written before generics, so Push and Pop use it.

`x.(T)` is a type assertion: "I know the box x holds a T, give me the T". If it holds something else, the program panics.

container/heap needs you to write Len, Less, Swap, Push and Pop, then call heap.Push(h, x), never h.Push(x) directly, or the heap order is not restored.

top_k_heap.go
package main
import (
"container/heap"
"fmt"
)
type wordCount struct {
word string
count int
}
// minHeap keeps the k best seen so far; the root is the weakest of them.
type minHeap []wordCount
func (h minHeap) Len() int { return len(h) }
func (h minHeap) Less(i, j int) bool {
if h[i].count != h[j].count {
return h[i].count < h[j].count
}
return h[i].word > h[j].word
}
func (h minHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] }
func (h *minHeap) Push(x any) { *h = append(*h, x.(wordCount)) } // x.(wordCount): "x holds a wordCount, give it to me"
func (h *minHeap) Pop() any {
old := *h
x := old[len(old)-1]
*h = old[:len(old)-1]
return x
}
// topKHeap: O(n + m log k). Worth it only when k is much smaller than m.
func topKHeap(counts map[string]int, k int) []wordCount {
h := &minHeap{}
for w, c := range counts {
heap.Push(h, wordCount{w, c})
if h.Len() > k {
heap.Pop(h) // drop the weakest
}
}
out := make([]wordCount, h.Len())
for i := len(out) - 1; i >= 0; i-- {
out[i] = heap.Pop(h).(wordCount)
}
return out
}
func main() {
counts := map[string]int{"go": 5, "rust": 2, "java": 4, "c": 1}
fmt.Println(topKHeap(counts, 2)) // [{go 5} {java 4}]
}
The short version
  • Map to count, slice to order. Maps never give you order.
  • Always tie-break, so the output is deterministic.
  • Heap of size k when k is much smaller than the number of distinct items.
Remember: If the output depends on map order, it is a bug that only fails sometimes.
Build this before moving on
Return the top K most frequent characters in a string. Decide whether you count bytes or runes, and say why.
Coding challenge · checked by tests
Write topK(words []string, k int) []string that returns the k most frequent words, most frequent first. Break ties by putting the word that sorts first (alphabetically) first, so the answer never depends on map order. If there are fewer than k distinct words, return them all. Example: [go is fun go is go], k = 2 → [go is].

Tab indents. Press Esc, then Tab, to leave the editor.

Then try these in your own editor
  1. 1The main task: Return the top K most frequent characters in a string. Decide whether you count bytes or runes, and say why.
  2. 2Remove the tie-break from the comparator and run the program ten times with words that tie. What happens?
  3. 3Group anagrams: ["eat","tea","tan","ate","nat"] → groups. Use a sorted-rune key in a map[string][]string.
  4. 4Dry run the sort on "b a b a c" with k = 2. Which tie-break decides the order?
  5. 5Return the k most frequent numbers in a []int using the same pattern.
Check yourselfTake the Go for coding rounds test