Go言語でPythonのcollections.Counterを簡単に実装してみた

Pythonのcollections.Counterは個人的にすごく好きなんですが，Go言語の練習を兼ねてGo言語で実装してみました．

完璧な実装ではなく，カウントできるのは文字列だけです．また，普通に実装間違えてる可能性もあります．自己責任でよろしくお願いします．

Pythonのcollections.Counterについてはここでも紹介しています．細かい仕様については公式ページ参照のこと．

以下コードです．

package main

import (
    "fmt"
    "sort"
)

type Counter struct {
    innerMap map[string]int
    sum      int
}

type KeyCountPair struct {
    key   string
    count int
}

type KeyCountPairList []KeyCountPair

func (kcpl KeyCountPairList) Len() int {
    return len(kcpl)
}

func (kcpl KeyCountPairList) Swap(i, j int) {
    kcpl[i], kcpl[j] = kcpl[j], kcpl[i]
}

func (kcpl KeyCountPairList) Less(i, j int) bool {
    return kcpl[i].count > kcpl[j].count
}

func (c *Counter) MostCommon(n int) KeyCountPairList {
    length := len(c.innerMap)
    pairs := make(KeyCountPairList, length)
    idx := 0
    for k, v := range c.innerMap {
        pairs[idx] = KeyCountPair{
            key:   k,
            count: v,
        }
        idx++
    }
    sort.Sort(pairs)
    return pairs[:n]
}

func (c *Counter) Elements() []string {
    if c.sum == 0 {
        for _, v := range c.innerMap {
            if v > 0 {
                c.sum += v
            }
        }
    }
    elements := make([]string, c.sum)
    idx := 0
    for k, v := range c.innerMap {
        if v > 0 {
            for i := 0; i < v; i++ {
                elements[idx] = k
                idx++
            }
        }
    }
    return elements
}

func (c *Counter) Subtract(tc *Counter) *Counter {
    for k := range c.innerMap {
        c.AddN(k, -tc.Count(k))
    }
    return c
}

func (c *Counter) Add(s string) {
    c.innerMap[s] += 1
    c.sum += 1
}

func (c *Counter) AddN(s string, n int) {
    c.innerMap[s] += n
    c.sum += n
}

func (c *Counter) Count(s string) int {
    return c.innerMap[s]
}

func NewCounter() *Counter {
    m := make(map[string]int)
    return &Counter{
        innerMap: m,
        sum:      0,
    }
}

func main() {
    c := NewCounter()
    c.AddN("test", 3)
    c.AddN("apolo", 4)
    c.AddN("google", 8)
    c.AddN("microsoft", 3)
    tc := NewCounter()
    tc.Add("test")
    fmt.Println(c)               // ==> &{map[microsoft:3 test:3 apolo:4 google:8] 18}
    fmt.Println(c.Count("word")) // ==> 0
    fmt.Println(c.Subtract(tc))  // ==> &{map[test:2 apolo:4 google:8 microsoft:3] 17}
    fmt.Println(c.Elements())    // ==> [test test apolo apolo apolo apolo google google google google google google google google microsoft microsoft microsoft]
    fmt.Println(c.MostCommon(2)) // ==> [{google 8} {apolo 4}]
}

sort.Sortで降順にソートする場合ってLessの不等号逆にするしかないってまじですか？