43 lines
616 B
Go
43 lines
616 B
Go
package collectionx
|
|
|
|
import "github.com/samber/lo"
|
|
|
|
type CountMapN[T comparable] struct {
|
|
N int
|
|
Map map[T][]int
|
|
}
|
|
|
|
func NewCountMapN[T comparable](n int) *CountMapN[T] {
|
|
return &CountMapN[T]{
|
|
N: n,
|
|
Map: make(map[T][]int),
|
|
}
|
|
}
|
|
|
|
func (cm *CountMapN[T]) Inc(key T, ith int, value int) {
|
|
m, ok := cm.Map[key]
|
|
if !ok {
|
|
m = make([]int, cm.N)
|
|
cm.Map[key] = m
|
|
}
|
|
|
|
m[ith] += value
|
|
}
|
|
|
|
func (cm *CountMapN[T]) Get(key T, ith int) int {
|
|
m, ok := cm.Map[key]
|
|
if !ok {
|
|
return 0
|
|
}
|
|
|
|
return m[ith]
|
|
}
|
|
|
|
func (cm *CountMapN[T]) Sum(key T) int {
|
|
m, ok := cm.Map[key]
|
|
if !ok {
|
|
return 0
|
|
}
|
|
return lo.Sum(m)
|
|
}
|