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.
- 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.
package mainimport ( "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.Comparewrites it for you. The oldersort.Slicetakes alessbool and still works. - `min` and `max` are built-ins since Go 1.21.
list[:min(k, len(list))]stops a panic whenkis 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.
package mainimport ( "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 []wordCountfunc (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}]}- 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.
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.
- 1The main task: Return the top K most frequent characters in a string. Decide whether you count bytes or runes, and say why.
- 2Remove the tie-break from the comparator and run the program ten times with words that tie. What happens?
- 3Group anagrams:
["eat","tea","tan","ate","nat"]→ groups. Use a sorted-rune key in amap[string][]string. - 4Dry run the sort on
"b a b a c"withk = 2. Which tie-break decides the order? - 5Return the
kmost frequent numbers in a[]intusing the same pattern.