如何用golang操作前缀树

98 阅读1分钟

前缀树定义

将字符串数组{“abc”,"abcd","abce","db"}加入前缀树

定义TrieNode

type TrieNode struct {
   pass int  
   end  int
   next []*TrieNode
}

pass 表示通过次数

end 表示以该字符结束的次数

添加abc时,将a添加得到

image.png

a的通过次数p=1

添加b得到

image.png

b的通过次数p=1

添加c得到

image.png

将c的end++

添加第二个字符串abcd得到

image.png

添加第三个字符串abce得到

image.png

添加第四个字符串db得到

image.png

我们可以很轻易知道 abc 出现的次数就是c的end = 1

以abc为前缀出现的次数就是c的pass = 3

删除一个元素时,需要先判断是否存在,再沿途将pass-1,最后end-1即可

package main

import "fmt"

type TrieNode struct {
   pass int
   end  int
   next []*TrieNode
}

func add(node *TrieNode, word string) {
   node.pass++
   for _, w := range word {
      index := w - 'a'
      if node.next[index] == nil {
         node.next[index] = &TrieNode{
            pass: 0,
            end:  0,
            next: make([]*TrieNode, 26),
         }
      }
      node = node.next[index]
      node.pass++
   }
   node.end++
}

func search(node *TrieNode, word string) int {
   for _, w := range word {
      index := w - 'a'
      if node.next[index] == nil {
         return 0
      }
      node = node.next[index]
   }
   return node.end
}

func prefixNumber(node *TrieNode, word string) int {
   for _, w := range word {
      index := w - 'a'
      if node.next[index] == nil {
         return 0
      }
      node = node.next[index]
   }
   return node.pass
}

func delete(node *TrieNode, word string) {
   if search(node, word) == 0 {
      return
   }
   node.pass--
   for _, w := range word {
      index := w - 'a'
      node.next[index].pass--
      if node.next[index].pass == 0 {
         node.next[index] = nil
         return
      }
      node = node.next[index]
   }
   node.end--
}

func main() {
   root := &TrieNode{next: make([]*TrieNode, 26)}
   add(root, "abc")
   add(root, "abcd")
   add(root, "abce")
   add(root, "bd")
   fmt.Println(search(root, "abc"))
   fmt.Println(prefixNumber(root, "abc"))
   fmt.Println(root.pass)

   delete(root, "abc")
   fmt.Println(search(root, "abc"))

   fmt.Println(root.pass)
}