Mimic Python's collection.Counter in Go

counter.go

package counter

import "container/heap"

// Implement the maxheap used in ic.MostCommon()

// KeyValueHeap is a nested slice that implements the heap interface
// It roughly resembles an ordered mapping where each element is a
// length-2 slice with keys as the 0th element and values as the 1st
// See https://golang.org/pkg/container/heap/ IntHeap example
type KeyValueHeap [][]int

func (h KeyValueHeap) Len() int { return len(h) }

// Use the "opposite" sign because we want Pop to give highest element
// The key is the 1st value from each pair
func (h KeyValueHeap) Less(i, j int) bool  { return h[i][1] < h[j][1] }
func (h KeyValueHeap) Swap(i, j int)       { h[i], h[j] = h[j], h[i] }
func (h *KeyValueHeap) Push(x interface{}) { *h = append(*h, x.([]int)) }
func (h *KeyValueHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

// Counter tallies counts of elements from a sequence and
// has methods that return maps with elements as keys and counts as values
type Counter interface {
	Count()
	MostCommon()
}

type IntCounter struct {
	a []int
}

// Counts returns a map of the elements to their corresponding frequency counts
func (ic IntCounter) Counts() map[int]int {
	counts := make(map[int]int)
	for _, i := range ic.a {
		counts[i]++
	}
	return counts
}

// MostCommon returns a map of the most common `n` elements to their counts
func (ic IntCounter) MostCommon(n int) map[int]int {
	h := &KeyValueHeap{}
	heap.Init(h)
	counts := ic.Counts()

	for k, v := range counts {
		pair := []int{k, v}
		heap.Push(h, pair)
	}

	result := make(map[int]int)
	// Pop the n largest results
	for i := 0; i < n; i++ {
		newpair := h.Pop().([]int)
		result[newpair[0]] = newpair[1]
	}
	return result
}