leetcode-go/leetcode/medium/677.map-sum-pairs.go

type trie struct {
	freq int
	word bool
	next [26]*trie // '26' is quite important!
}

func (t *trie) insert(key string, val int) {
	res := t.search(key)
	if res != nil && res.word {
		val -= res.freq
	}
	for _, r := range key {
		i := r - 'a'
		if t.next[i] == nil {
			t.next[i] = &trie{}
		}
		t = t.next[i]
		t.freq += val
	}
	t.word = true
}

func (t *trie) search(key string) *trie {
	for _, r := range key {
		t = t.next[r-'a']
		if t == nil {
			break
		}
	}
	return t
}

type MapSum struct {
	tree trie
}

/** Initialize your data structure here. */
func Constructor() MapSum {
	return MapSum{}
}

func (this *MapSum) Insert(key string, val int) {
	this.tree.insert(key, val)
}

func (this *MapSum) Sum(prefix string) int {
	res := this.tree.search(prefix)
	if res == nil {
		return 0
	}
	return res.freq
}

/**
 * Your MapSum object will be instantiated and called as such:
 * obj := Constructor();
 * obj.Insert(key,val);
 * param_2 := obj.Sum(prefix);
 */