trie_v1.go

package phrasetrie

type node struct {
	children map[rune]*node
	value    string
}

func newNode() *node {
	return &node{
		children: make(map[rune]*node),
	}
}

type Trie struct {
	root *node
}

func NewTrie(keys ...string) *Trie {
	t := &Trie{
		root: newNode(),
	}
	for _, k := range keys {
		t.add(k)
	}
	return t
}

func (t *Trie) add(key string) {
	n := t.root
	for _, c := range key {
		if _, ok := n.children[c]; !ok {
			n.children[c] = newNode()
		}
		n = n.children[c]
	}
	n.value = key
}

func (t *Trie) Match(word string) bool {
	_, idx, _ := t.search([]rune(word), 0, false)
	return 0 <= idx
}

func (t *Trie) FindString(word string) string {
	v, _, _ := t.search([]rune(word), 0, false)
	return v
}

func (t *Trie) FindLongestString(word string) string {
	v, _, _ := t.search([]rune(word), 0, true)
	return v
}

func (t *Trie) ReplaceAll(word, repl string) string {
	return t.ReplaceAllFunc(word, func(_ string) string { return repl })
}

func (t *Trie) ReplaceAllFunc(word string, repl func(string) string) string {
	chars := []rune(word)
	dist := make([]rune, 0, len(chars))
	start := 0
	for {
		found, idx, end := t.search(chars, start, true)
		if idx < 0 {
			dist = append(dist, chars[start:]...)
			break
		}
		dist = append(dist, chars[start:idx]...)
		dist = append(dist, []rune(repl(found))...)
		start = end + 1
	}
	return string(dist)
}

func (t *Trie) search(chars []rune, start int, longest bool) (string, int, int) {
	end := len(chars)
	var v string
	var e int
	for i := start; i < end; i++ {
		if n, ok := t.root.children[chars[i]]; ok {
			for j := i + 1; j < end; j++ {
				if n, ok = n.children[chars[j]]; ok {
					if n.value != "" {
						v = n.value
						e = j
						if !longest {
							break
						}
					}
				} else {
					break
				}
			}
			if v != "" {
				return v, i, e
			}
		}
	}
	return "", -1, -1
}

goos: linux
goarch: amd64
BenchmarkReplace_30/regexp-2                2618            385521 ns/op            5392 B/op          3 allocs/op
BenchmarkReplace_30/Replacer-2             57013             20246 ns/op            5408 B/op          3 allocs/op
BenchmarkReplace_30/Trie-2                 23221             58154 ns/op            9600 B/op          3 allocs/op
BenchmarkReplace_Many/regexp-2               481           2238582 ns/op           10081 B/op         10 allocs/op
BenchmarkReplace_Many/Replacer-2           42528             31628 ns/op            5031 B/op          3 allocs/op
BenchmarkReplace_Many/Trie-2               14562             72307 ns/op           16000 B/op          4 allocs/op