...
1
16
17 package util
18
19
20 type Trie struct {
21 children map[byte]*Trie
22 wordTail bool
23 word string
24 }
25
26
27 func NewTrie(list []string) Trie {
28 ret := Trie{
29 children: make(map[byte]*Trie),
30 wordTail: false,
31 }
32 for _, v := range list {
33 ret.Add(v)
34 }
35 return ret
36 }
37
38
39 func (t *Trie) Add(v string) {
40 root := t
41 for _, b := range []byte(v) {
42 child, exists := root.children[b]
43 if !exists {
44 child = &Trie{
45 children: make(map[byte]*Trie),
46 wordTail: false,
47 }
48 root.children[b] = child
49 }
50 root = child
51 }
52 root.wordTail = true
53 root.word = v
54 }
55
56
57 func (t *Trie) HasPrefix(v string) bool {
58 _, has := t.GetPrefix(v)
59 return has
60 }
61
62
63 func (t *Trie) GetPrefix(v string) (string, bool) {
64 root := t
65 if root.wordTail {
66 return root.word, true
67 }
68 for _, b := range []byte(v) {
69 child, exists := root.children[b]
70 if !exists {
71 return "", false
72 }
73 if child.wordTail {
74 return child.word, true
75 }
76 root = child
77 }
78 return "", false
79 }
80
View as plain text