-
Notifications
You must be signed in to change notification settings - Fork 1
/
rtrie.go
59 lines (51 loc) · 979 Bytes
/
rtrie.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
package utils
type TrieNode struct {
Value interface{}
Children []*TrieNode
}
const (
defaultR = 26
)
type RTrie struct {
Root *TrieNode
r int
}
func NewRTrie(r int) *RTrie {
if r <= 0 {
r = defaultR
}
return &RTrie{Root: &TrieNode{Children: make([]*TrieNode, r)}, r: r}
}
func (rt *RTrie) Put(key string, value interface{}) {
p := rt.Root
for _, c := range key {
idx := c - 'a'
if p.Children[idx] == nil {
p.Children[idx] = &TrieNode{
Children: make([]*TrieNode, rt.r),
}
}
p = p.Children[idx]
}
p.Value = value
}
func (rt *RTrie) Contains(query string) bool {
return rt.Get(query) != nil
}
func (rt *RTrie) Get(query string) interface{} {
node := rget(rt.Root, query, 0)
if node == nil {
return nil
}
return node.Value
}
func rget(node *TrieNode, query string, d int) *TrieNode {
if node == nil {
return node
}
if d == len(query) {
return node
}
idx := query[d] - 'a'
return rget(node.Children[idx], query, d+1)
}