148 lines
2.5 KiB
Go
148 lines
2.5 KiB
Go
package trie
|
|
|
|
import (
|
|
"fmt"
|
|
)
|
|
|
|
type (
|
|
Trie struct {
|
|
root *node
|
|
size int
|
|
}
|
|
node struct {
|
|
key interface{}
|
|
value interface{}
|
|
next [256]*node
|
|
}
|
|
iterator struct {
|
|
step int
|
|
node *node
|
|
prev *iterator
|
|
}
|
|
)
|
|
|
|
func toBytes(obj interface{}) []byte {
|
|
switch o := obj.(type) {
|
|
case []byte:
|
|
return o
|
|
case string:
|
|
return []byte(o)
|
|
}
|
|
return []byte(fmt.Sprint(obj))
|
|
}
|
|
|
|
func New() *Trie {
|
|
return &Trie{nil,0}
|
|
}
|
|
func (this *Trie) Do(handler func(interface{},interface{})bool) {
|
|
if this.size > 0 {
|
|
this.root.do(handler)
|
|
}
|
|
}
|
|
func (this *Trie) Get(key interface{}) interface{} {
|
|
if this.size == 0 {
|
|
return nil
|
|
}
|
|
|
|
bs := toBytes(key)
|
|
cur := this.root
|
|
for i := 0; i < len(bs); i++ {
|
|
if cur.next[bs[i]] != nil {
|
|
cur = cur.next[bs[i]]
|
|
} else {
|
|
return nil
|
|
}
|
|
}
|
|
return cur.value
|
|
}
|
|
func (this *Trie) Has(key interface{}) bool {
|
|
return this.Get(key) != nil
|
|
}
|
|
func (this *Trie) Init() {
|
|
this.root = nil
|
|
this.size = 0
|
|
}
|
|
func (this *Trie) Insert(key interface{}, value interface{}) {
|
|
if this.size == 0 {
|
|
this.root = newNode()
|
|
}
|
|
|
|
bs := toBytes(key)
|
|
cur := this.root
|
|
for i := 0; i < len(bs); i++ {
|
|
if cur.next[bs[i]] != nil {
|
|
cur = cur.next[bs[i]]
|
|
} else {
|
|
cur.next[bs[i]] = newNode()
|
|
cur = cur.next[bs[i]]
|
|
}
|
|
}
|
|
if cur.key == nil {
|
|
this.size++
|
|
}
|
|
cur.key = key
|
|
cur.value = value
|
|
}
|
|
func (this *Trie) Len() int {
|
|
return this.size
|
|
}
|
|
func (this *Trie) Remove(key interface{}) interface{} {
|
|
if this.size == 0 {
|
|
return nil
|
|
}
|
|
bs := toBytes(key)
|
|
cur := this.root
|
|
|
|
for i := 0; i < len(bs); i++ {
|
|
if cur.next[bs[i]] != nil {
|
|
cur = cur.next[bs[i]]
|
|
} else {
|
|
return nil
|
|
}
|
|
}
|
|
|
|
// TODO: cleanup dead nodes
|
|
|
|
val := cur.value
|
|
|
|
if cur.value != nil {
|
|
this.size--
|
|
cur.value = nil
|
|
cur.key = nil
|
|
}
|
|
return val
|
|
}
|
|
func (this *Trie) String() string {
|
|
str := "{"
|
|
i := 0
|
|
this.Do(func(k, v interface{}) bool {
|
|
if i > 0 {
|
|
str += ", "
|
|
}
|
|
str += fmt.Sprint(k, ":", v)
|
|
i++
|
|
return true
|
|
})
|
|
str += "}"
|
|
return str
|
|
}
|
|
|
|
func newNode() *node {
|
|
var next [256]*node
|
|
return &node{nil,nil,next}
|
|
}
|
|
func (this *node) do(handler func(interface{}, interface{}) bool) bool {
|
|
for i := 0; i < 256; i++ {
|
|
if this.next[i] != nil {
|
|
if this.next[i].key != nil {
|
|
if !handler(this.next[i].key, this.next[i].value) {
|
|
return false
|
|
}
|
|
}
|
|
if !this.next[i].do(handler) {
|
|
return false
|
|
}
|
|
}
|
|
}
|
|
return true
|
|
} |