day 232: prefix map sum trie

day232/problem.go

@@ -0,0 +1,46 @@
+package day232
+
+// PrefixMapSum represents a prefix tree collection of sums.
+type PrefixMapSum interface {
+	Insert(key string, value int)
+	Sum(prefix string) int
+}
+
+type prefixMapSumTrie struct {
+	sum  int
+	next map[rune]*prefixMapSumTrie
+}
+
+// NewPrefixMapSum returns a new instance.
+func NewPrefixMapSum() PrefixMapSum {
+	return &prefixMapSumTrie{}
+}
+
+func (pms *prefixMapSumTrie) Insert(key string, value int) {
+	curr := pms
+	for _, r := range key {
+		if curr.next == nil {
+			curr.next = make(map[rune]*prefixMapSumTrie)
+		}
+		if _, found := curr.next[r]; !found {
+			curr.next[r] = &prefixMapSumTrie{}
+		}
+		curr.next[r].sum += value
+		curr = curr.next[r]
+	}
+}
+
+func (pms *prefixMapSumTrie) Sum(prefix string) int {
+	curr := pms
+	for _, r := range prefix {
+		if curr.next == nil || curr.next[r] == nil {
+			curr = nil
+			break
+		}
+		curr = curr.next[r]
+	}
+	if curr != nil {
+		return curr.sum
+	}
+	return 0
+}

day232/problem_test.go

@@ -0,0 +1,33 @@
+package day232
+
+import "testing"
+
+func TestPrefixMapSum(t *testing.T) {
+	t.Parallel()
+	pms := NewPrefixMapSum()
+	pms.Insert("columnar", 3)
+	if sum := pms.Sum("col"); sum != 3 {
+		t.Errorf("Sum should be 3 here got %d", sum)
+	}
+	pms.Insert("column", 2)
+	if sum := pms.Sum("col"); sum != 5 {
+		t.Errorf("Sum should be 5 here got %d", sum)
+	}
+	if sum := pms.Sum("columns"); sum != 0 {
+		t.Errorf("Sum should be 0 here got %d", sum)
+	}
+	if sum := pms.Sum("hi there"); sum != 0 {
+		t.Errorf("Sum should be 0 here got %d", sum)
+	}
+}
+
+func BenchmarkPrefixMapSum(b *testing.B) {
+	pms := NewPrefixMapSum()
+	for i := 0; i < b.N; i++ {
+		pms.Insert("columnar", 3)
+		pms.Insert("column", 2)
+		pms.Sum("col")
+		pms.Sum("columns")
+		pms.Sum("hi there")
+	}
+}