Go实现的前缀树Trie

实现trie的增删查操作

package trie

type Node struct {
	Character, Word string
	IsEnd           bool
	Children        map[string]*Node
	Count           int
}

type children map[string]*Node

type Trie struct {
	Root     *Node
	alphabet []string
}

func (t *Trie) newChildren() children {
	c := children{}
	for _, char := range t.alphabet {
		c[char] = nil
	}
	return c
}

func (t *Trie) NewNode(existingSeq, character string, isEnd bool) *Node {
	return &Node{
		Character: character,
		Word:      existingSeq + character,
		IsEnd:     isEnd,
		Children:  t.newChildren(),
		Count:     0,
	}
}

func NewTrie(alphabet []string) *Trie {
	t := &Trie{}
	t.alphabet = alphabet
	t.Root = t.NewNode(``, ``, true)
	return t
}

func (n *Node) GetChild(c string) *Node {
	child := n.Children[c]
	return child
}

func (t *Trie) InsertWord(word string) bool {
	if t.Root.IsEnd {
		t.Root.IsEnd = false
	}
	if t.SearchWord(word) {
		return false
	}
	cur := t.Root
	for index, c := range word {
		cur.Count++
		if cur.GetChild(string(c)) != nil {
			cur = cur.Children[string(c)]
			continue
		}
		newNode := t.NewNode(word[:index], string(c), false)
		cur.Children[string(c)] = newNode
		cur = cur.GetChild(string(c))
	}
	cur.Count++
	cur.IsEnd = true
	return true
}

func (t *Trie) IsPrefixExist(prefix string) bool {
	cur := t.Root
	for _, c := range prefix {
		cur = cur.GetChild(string(c))
		if cur == nil {
			return false
		}
	}
	return true
}

func (t *Trie) SearchWord(word string) bool {
	cur := t.Root
	for _, c := range word {
		cur = cur.GetChild(string(c))
		if cur == nil {
			return false
		}
	}
	return cur.IsEnd
}

// DeleteWord returns false when the word doesn't exist in trie
func (t *Trie) DeleteWord(word string) bool {
	cur := t.Root
	prev := cur
	pastNodes := []*Node{cur}
	for _, c := range word {
		cur = cur.GetChild(string(c))
		pastNodes = append(pastNodes, cur)
		if cur == nil {
			return false
		}
		if cur.Count == 1 {
			// update count for all past nodes
			for _, node := range pastNodes[:len(pastNodes)-1] {
				node.Count--
			}
			prev.Children[string(c)] = nil
			return true
		}
		prev = cur
	}
	if cur.IsEnd {
		for _, node := range pastNodes {
			node.Count--
		}
		cur.IsEnd = false
		return true
	}
	return false
}

func (t *Trie) findWords(node *Node, words *[]string) {
	if node.IsEnd {
		*words = append(*words, node.Word)
	}
	for _, c := range t.alphabet {
		cur := node.GetChild(c)
		if cur != nil {
			t.findWords(cur, words)
		}
	}
}

func (t *Trie) FindWords() []string {
	words := []string{}
	t.findWords(t.Root, &words)
	return words
}