My-Notes

Search

SearchSearch

026.前缀树原理和代码详解

Dec 26, 2023, 8 min read

前置知识：知道什么是树结构，比如二叉树：数据结构分类、二叉树基本概念；知道为什么推荐静态数组的方式实现各种结构：推荐静态空间的实现；知道哈希表怎么用：哈希表的使用

前缀树

前缀树又叫字典树 ，英文名 trie

每个样本都从头节点开始，根据前缀字符或者前缀数字建出来的一棵大树

没有路就新建节点；已经有路了，就复用节点

前缀树的使用场景：需要根据前缀信息来查询的场景

前缀树的优点：根据前缀信息选择树上的分支，可以节省大量的时间

前缀树的缺点：比较浪费空间，和总字符数量和字符的种类有关

前缀树一般把字符放在路上，节点放置定制信息，如：pass、end 等

这节课是前缀树的原理和代码讲解。下节课是前缀树相关题目讲解，展示前缀树在解题时的常见用法。

实现前缀树

设计

Trie()                          // 构造方法
void insert(String word)        // 将字符串 word 插入前缀树中
int search(String word)         // 返回前缀树中字符串 word 的实例个数
int prefixNumber(String prefix) // 返回前缀树中以 prefix 为前缀的字符串个数
void delete(String word)        // 从前缀树中移除字符串 word（移除 1 个）

测试链接

类描述的实现方式

虽然最常用，但是是动态结构，刷题时不推荐

package main
 
import (
    "bufio"
    "fmt"
    "os"
    "strconv"
)
 
func main() {
    in := bufio.NewScanner(os.Stdin)
    in.Split(bufio.ScanWords)
    out := bufio.NewWriterSize(os.Stdout, 4096)
    for in.Scan() {
        n, _ := strconv.Atoi(in.Text())
        if n <= 0 {
            continue
        }
        trie := NewTrie()
        for ; n > 0; n-- {
            in.Scan()
            op := in.Text()
            in.Scan()
            word := in.Text()
            switch op {
            case "1":
                trie.Insert(word)
            case "2":
                trie.Delete(word)
            case "3":
                if trie.Search(word) > 0 {
                    fmt.Fprintln(out, "YES")
                } else {
                    fmt.Fprintln(out, "NO")
                }
            case "4":
                fmt.Fprintln(out, trie.PrefixNumber(word))
            }
        }
    }
    out.Flush()
}
 
type TrieNode struct {
    pass  int           // 经过节点结尾
    end   int           // 以当前节点结尾
    nexts [26]*TrieNode // 下一个节点，a-z
}
 
func NewTrie() *TrieNode {
    return new(TrieNode)
}
 
// 将字符串 word 插入前缀树中
func (t *TrieNode) Insert(word string) {
    n := len(word)
    node := t
    node.pass++
    for i := 0; i < n; i++ {
        idx := word[i] - 'a'
        if node.nexts[idx] == nil {
            node.nexts[idx] = new(TrieNode) // 没有路就新建节点
        }
        node = node.nexts[idx]
        node.pass++
    }
    node.end++
}
 
// 返回前缀树中字符串 word 的实例个数
func (t *TrieNode) Search(word string) int {
    n := len(word)
    node := t
    for i := 0; i < n; i++ {
        node = node.nexts[word[i]-'a']
        if node == nil {
            return 0
        }
    }
    return node.end
}
 
// 返回前缀树中以 prefix 为前缀的字符串个数
func (t *TrieNode) PrefixNumber(prefix string) int {
    n := len(prefix)
    node := t
    for i := 0; i < n; i++ {
        node = node.nexts[prefix[i]-'a']
        if node == nil {
            return 0
        }
    }
    return node.pass
}
 
// 从前缀树中移除字符串 word（移除 1 个）
func (t *TrieNode) Delete(word string) {
    if t.Search(word) == 0 {
        return
    }
    n := len(word)
    node := t
    node.pass--
    for i := 0; i < n; i++ {
        idx := word[i] - 'a'
        node.nexts[idx].pass--
        if node.nexts[idx].pass == 0 {
            // 节省空间，也为了 GC
            node.nexts[idx] = nil
            return
        }
        node = node.nexts[idx]
    }
    node.end--
}

本题中，路为 a-z，路的可能性范围较小，用固定数组实现路：索引快、空间占用可接受

路的可能性范围较大时，用哈希表实现路：空间在需要时才开辟

静态数组的实现方式

推荐！不仅笔试，就连比赛也能保证使用

package main
 
import (
	"bufio"
	"fmt"
	"os"
	"strconv"
)
 
const (
	MAXN = 150001
)
 
var (
	// 四个变量实现 Trie
	tree = [MAXN][26]int{}
	pass = [MAXN]int{}
	end  = [MAXN]int{}
	cnt  = 1
 
	n    int
	op   string
	word string
)
 
func main() {
	in := bufio.NewScanner(os.Stdin)
	in.Split(bufio.ScanWords)
	out := bufio.NewWriterSize(os.Stdout, 4096)
	for in.Scan() {
		n, _ = strconv.Atoi(in.Text())
		if n <= 0 {
			continue
		}
		for ; n > 0; n-- {
			in.Scan()
			op = in.Text()
			in.Scan()
			word = in.Text()
			switch op {
			case "1":
				Insert(word)
			case "2":
				Delete(word)
			case "3":
				if Search(word) > 0 {
					fmt.Fprintln(out, "YES")
				} else {
					fmt.Fprintln(out, "NO")
				}
			case "4":
				fmt.Fprintln(out, PrefixNumber(word))
			}
		}
		Clear()
	}
	out.Flush()
}
 
// 将字符串 word 插入前缀树中
func Insert(word string) {
	n := len(word)
	treeIdx := 0
	pass[treeIdx]++
	for i := 0; i < n; i++ {
		idx := word[i] - 'a'
		if tree[treeIdx][idx] == 0 {
			tree[treeIdx][idx] = cnt // 没有路就新建节点
			cnt++
		}
		treeIdx = tree[treeIdx][idx]
		pass[treeIdx]++
	}
	end[treeIdx]++
}
 
// 返回前缀树中字符串 word 的实例个数
func Search(word string) int {
	n := len(word)
	treeIdx := 0
	for i := 0; i < n; i++ {
		treeIdx = tree[treeIdx][word[i]-'a']
		if treeIdx == 0 {
			return 0
		}
	}
	return end[treeIdx]
}
 
// 返回前缀树中以 prefix 为前缀的字符串个数
func PrefixNumber(prefix string) int {
	n := len(prefix)
	treeIdx := 0
	for i := 0; i < n; i++ {
		treeIdx = tree[treeIdx][prefix[i]-'a']
		if treeIdx == 0 {
			return 0
		}
	}
	return pass[treeIdx]
}
 
// 从前缀树中移除字符串 word（移除 1 个）
func Delete(word string) {
	if Search(word) == 0 {
		return
	}
	n := len(word)
	treeIdx := 0
	pass[treeIdx]--
	for i := 0; i < n; i++ {
		idx := word[i] - 'a'
		pass[tree[treeIdx][idx]]--
		if pass[tree[treeIdx][idx]] == 0 {
			tree[treeIdx][idx] = 0
			return
		}
		treeIdx = tree[treeIdx][idx]
	}
	end[treeIdx]--
}
 
// 清除 Trie 树的信息，方便下次使用
func Clear() {
	for i := 0; i < cnt; i++ {
		for j := 0; j < 26; j++ {
			tree[i][j] = 0
		}
		pass[i] = 0
		end[i] = 0
	}
	cnt = 1
}

一切都是静态数组来实现，提交准备好够用的空间（MAXN 要足够：最多是 字符串长度 * 字符串数量，但是前缀树又有大量的空间压缩，所以一般情况下要远远小于这个数，可以一开始给一个够用的，然后逐渐减少量，试出来一个合适的）

如果路的可能性范围较大，因为不用动态结构就不能用哈希表了。需要用每一位的信息建树，如：第 1789 条路建树为 1->7->8->9->#，相当于 tree[MAXN][11]，当然也可以用其他进制。这样的话 MAXN 也要适量增大（这道题）

Graph View

Backlinks