-
Notifications
You must be signed in to change notification settings - Fork 0
/
trie.go
118 lines (96 loc) · 1.97 KB
/
trie.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
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
package trie
import "sync"
type Trie struct {
sync.RWMutex
Root *Node
}
type Node struct {
Symbol rune
Parent *Node
Children map[rune]*Node
Data interface{}
}
func NewTrie() *Trie {
return &Trie{
Root: &Node{
Children: make(map[rune]*Node),
},
}
}
func (trie *Trie) Insert(key string, data interface{}) {
trie.Lock()
node := trie.Root
for _, r := range key {
childNode := node.Children[r]
if childNode == nil {
childNode = &Node{
Symbol: r,
Parent: node,
Children: make(map[rune]*Node),
}
node.Children[r] = childNode
}
node = childNode
}
node.Data = data
trie.Unlock()
}
func (trie *Trie) Search(key string) interface{} {
trie.RLock()
node := trie.Root.findNode(key)
trie.RUnlock()
if node == nil {
return nil
}
return node.Data
}
func (trie *Trie) HasPrefix(prefix string) map[string]interface{} {
var results = make(map[string]interface{})
trie.RLock()
defer trie.RUnlock()
node := trie.Root.findNode(prefix)
if node == nil {
return results
}
if node.Data != nil {
results[prefix] = node.Data
}
var findResults func(*Node, string)
findResults = func(node *Node, prefix string) {
for r, childNode := range node.Children {
childPrefix := prefix + string(r)
if childNode.Data != nil {
results[childPrefix] = childNode.Data
}
findResults(childNode, childPrefix)
}
}
findResults(node, prefix)
return results
}
func (trie *Trie) Delete(key string) bool {
trie.Lock()
defer trie.Unlock()
node := trie.Root.findNode(key)
if node == nil || node.Data == nil {
return false
}
node.Data = nil
for node.Data == nil && len(node.Children) == 0 && node.Parent != nil {
delete(node.Parent.Children, node.Symbol)
parent := node.Parent
node.Parent = nil
node = parent
}
return true
}
// Ensure it is called inside the mutex lock
func (node *Node) findNode(key string) *Node {
for _, r := range key {
node = node.Children[r]
if node == nil {
return nil
}
}
return node
}