跳转至

LC208 - 实现 Trie (前缀树)

ImplementTrie.go
package Graph

type TrieNode struct {
    pass, end int
    nexts     [26]*TrieNode
}
type Trie struct {
    root *TrieNode
}

func Constructor() Trie {
    trie := new(Trie)
    trie.root = new(TrieNode)
    return *trie
}

func (this *Trie) Insert(word string) {
    cur := this.root
    cur.pass++
    for _, ch := range word {
        if cur.nexts[ch-'a'] == nil {
            cur.nexts[ch-'a'] = new(TrieNode)
        }
        cur = cur.nexts[ch-'a']
        cur.pass++
    }
    cur.end++
}

func (this *Trie) Search(word string) bool {
    cur := this.root
    for _, ch := range word {
        if cur.nexts[ch-'a'] == nil {
            return false
        }
        cur = cur.nexts[ch-'a']
    }
    return cur.end > 0
}

func (this *Trie) StartsWith(prefix string) bool {
    cur := this.root
    for _, ch := range prefix {
        if cur.nexts[ch-'a'] == nil {
            return false
        }
        cur = cur.nexts[ch-'a']
    }
    return true
}

func (this *Trie) Count(word string) int {
    cur := this.root
    for _, ch := range word {
        if cur.nexts[ch-'a'] == nil {
            return 0
        }
        cur = cur.nexts[ch-'a']
    }
    return cur.end
}

func (this *Trie) CountPrefix(prefix string) int {
    cur := this.root
    for _, ch := range prefix {
        if cur.nexts[ch-'a'] == nil {
            return 0
        }
        cur = cur.nexts[ch-'a']
    }
    return cur.pass
}

func (this *Trie) Delete(word string) {
    if !this.Search(word) {
        return
    }
    cur := this.root
    cur.pass--
    for _, ch := range word {
        cur.nexts[ch-'a'].pass--
        if cur.nexts[ch-'a'].pass == 0 {
            cur.nexts[ch-'a'] = nil
            return
        }
        cur = cur.nexts[ch-'a']
    }
    cur.end--
}