focus/tree/tried/tried.go

199 lines
3.6 KiB
Go
Raw Normal View History

package tried
import "github.com/davecgh/go-spew/spew"
2019-08-19 07:24:22 +00:00
// func (ts TriedString) WordIndex(idx uint) uint {
// w := ts[idx]
// if w >= 'a' && w <= 'z' {
// return uint(w) - 'a'
// } else if w >= 'A' && w <= 'Z' {
// return uint(w) - 'A' + 26
// } else {
// return uint(w) - '0' + 52
// }
// }
type Tried struct {
root *Node
wiStore *wordIndexStore
}
type Node struct {
data []*Node
value interface{}
}
// New 默认 WordIndexLower 意味着只支持小写
func New() *Tried {
tried := &Tried{}
tried.root = new(Node)
tried.wiStore = WordIndexDict[WordIndexLower]
return tried
}
// NewWithWordType 选择单词的类型 WordIndexLower 意味着只支持小写
func NewWithWordType(t WordIndexType) *Tried {
tried := &Tried{}
tried.root = new(Node)
tried.wiStore = WordIndexDict[t]
return tried
2019-08-19 06:30:24 +00:00
}
2019-08-19 07:28:09 +00:00
func (tried *Tried) Put(words string, values ...interface{}) {
cur := tried.root
var n *Node
bytes := []byte(words)
for i := 0; i < len(bytes); i++ {
w := tried.wiStore.Byte2Index(bytes[i])
if cur.data == nil {
cur.data = make([]*Node, tried.wiStore.DataSize)
}
if n = cur.data[w]; n == nil {
n = new(Node)
cur.data[w] = n
}
cur = n
}
if values != nil {
cur.value = values[0]
} else {
cur.value = tried
}
}
2019-08-19 07:28:09 +00:00
func (tried *Tried) Get(words string) interface{} {
cur := tried.root
var n *Node
bytes := []byte(words)
for i := 0; i < len(bytes); i++ {
w := tried.wiStore.Byte2Index(bytes[i]) //TODO: 升级Index 函数
if n = cur.data[w]; n == nil {
return nil
}
cur = n
}
return n.value
}
2019-08-19 07:28:09 +00:00
func (tried *Tried) Has(words string) bool {
return tried.Get(words) != nil
}
func (tried *Tried) HasPrefix(words string) bool {
cur := tried.root
var n *Node
bytes := []byte(words)
for i := 0; i < len(bytes); i++ {
w := tried.wiStore.Byte2Index(bytes[i]) //TODO: 升级Index 函数
if n = cur.data[w]; n == nil {
return false
}
cur = n
}
return true
}
func (tried *Tried) PrefixWords(words string) []string {
cur := tried.root
var n *Node
bytes := []byte(words)
var header []byte
for i := 0; i < len(bytes); i++ {
curbyte := bytes[i]
header = append(header, curbyte)
w := tried.wiStore.Byte2Index(curbyte)
if n = cur.data[w]; n == nil {
return nil
}
cur = n
}
var result []string
var traversal func([]byte, *Node)
traversal = func(prefix []byte, cur *Node) {
for i, n := range cur.data {
if n != nil {
nextPrefix := append(prefix, tried.wiStore.Index2Byte(uint(i)))
traversal(nextPrefix, n)
if n.value != nil {
result = append(result, string(append(header, nextPrefix...)))
}
}
}
}
// 拼接头
if n != nil {
if n.value != nil {
result = append(result, string(header))
}
traversal([]byte{}, n)
}
return result
}
func (tried *Tried) Traversal(every func(cidx uint, value interface{}) bool) {
var traversal func(*Node)
traversal = func(cur *Node) {
if cur != nil {
for i, n := range cur.data {
if n != nil {
if n.value != nil {
if !every(uint(i), n.value) {
return
}
}
traversal(n)
}
}
}
}
root := tried.root
traversal(root)
}
func (tried *Tried) WordsArray() []string {
var result []string
var traversal func([]byte, *Node)
traversal = func(prefix []byte, cur *Node) {
for i, n := range cur.data {
if n != nil {
nextPrefix := append(prefix, tried.wiStore.Index2Byte(uint(i)))
2019-08-20 03:11:10 +00:00
traversal(nextPrefix, n)
if n.value != nil {
2019-08-20 03:11:10 +00:00
result = append(result, string(nextPrefix))
}
}
}
}
if tried.root != nil {
traversal([]byte{}, tried.root)
}
return result
}
func (tried *Tried) String() string {
return spew.Sprint(tried.WordsArray())
}