github.com/qiuhoude/go-web@v0.0.0-20220223060959-ab545e78f20d/algorithm/datastructures/tree/trie/trie.go (about)

     1  // 前缀树
     2  package trie
     3  
     4  type Trie struct {
     5  	root *node
     6  	size int // 存储词的数量
     7  }
     8  
     9  type node struct {
    10  	children map[rune]*node // 孩子节点
    11  	isEnd    bool           // 结束标识符
    12  	data     interface{}    // 存储的数据
    13  }
    14  
    15  func NewTrie() *Trie {
    16  	n := newNode()
    17  	return &Trie{root: n, size: 0}
    18  }
    19  
    20  func newNode() *node {
    21  	n := new(node)
    22  	n.children = make(map[rune]*node)
    23  	return n
    24  }
    25  
    26  func (t Trie) Size() int {
    27  	return t.size
    28  }
    29  
    30  // 添加
    31  func (t *Trie) Add(w string, data interface{}) {
    32  	cur := t.root
    33  	for _, v := range []rune(w) {
    34  
    35  		n, ok := cur.children[v]
    36  		if !ok {
    37  			n = newNode()
    38  			cur.children[v] = n
    39  		}
    40  		cur = n
    41  	}
    42  	cur.data = data
    43  	if !cur.isEnd {
    44  		cur.isEnd = true
    45  		t.size++
    46  	}
    47  }
    48  
    49  // 删除词
    50  func (t *Trie) Remove(w string) bool {
    51  	cur := t.root
    52  	var stack []*node
    53  	runeArr := []rune(w)
    54  	for _, v := range runeArr {
    55  
    56  		n, ok := cur.children[v]
    57  		if !ok { // 没有找到该单词不用删除
    58  			return false
    59  		}
    60  		stack = append(stack, n)
    61  		cur = n
    62  	}
    63  	t.size--
    64  	// 结尾标识改掉
    65  	cur.isEnd = false
    66  	if len(cur.children) == 1 { // 只有自己一个字符 没有后续的 就可以进行移除操作
    67  		for i := len(stack) - 1; i >= 0; i-- {
    68  			c := runeArr[i]
    69  			n := stack[i]
    70  			delete(n.children, c)
    71  			if cur.isEnd {
    72  				break
    73  			}
    74  		}
    75  	}
    76  	return true
    77  }
    78  
    79  func (t *Trie) Find(w string) interface{} {
    80  	cur := t.root
    81  	for _, v := range []rune(w) {
    82  
    83  		n, ok := cur.children[v]
    84  		if !ok {
    85  			return nil
    86  		}
    87  		cur = n
    88  	}
    89  	if !cur.isEnd { // 不是结束就返回nil
    90  		return nil
    91  	}
    92  	return cur.data
    93  }
    94  
    95  func (t *Trie) Contains(w string) bool {
    96  	return t.Find(w) != nil
    97  }
    98  
    99  // 前缀搜索
   100  func (t *Trie) SearchPrefix(prefix string) []interface{} {
   101  	cur := t.root
   102  	for _, v := range []rune(prefix) {
   103  		n, ok := cur.children[v]
   104  		if !ok {
   105  			return nil
   106  		}
   107  		cur = n
   108  	}
   109  	ret := make([]interface{}, 0, 8)
   110  	searchNode(cur, &ret)
   111  	return ret
   112  }
   113  
   114  func searchNode(n *node, ret *[]interface{}) {
   115  	if n.isEnd {
   116  		*ret = append(*ret, n.data)
   117  	}
   118  	if len(n.children) <= 0 {
   119  		return
   120  	}
   121  	for _, cn := range n.children {
   122  		searchNode(cn, ret)
   123  	}
   124  }