Golang高效前缀树(Prefix Tree)实现插件库ptrie的使用

Trie (前缀树)

前缀树是一种空间优化的树形数据结构，其中每个节点都与其父节点合并。与常规树不同，前缀树中每个节点的键是按块进行比较的。

前缀树有以下应用场景：

• 文本文档搜索
• 基于规则的匹配
• 为字符串键构建关联数组

字符比较复杂度：

• 暴力搜索：O(d n k)
• 前缀树：O(d log(k))

其中：

• d：文档中的字符数
• n：关键词数量
• k：平均关键词长度

使用示例

1. 基本操作

package main

import (
	"fmt"
	"log"
	"github.com/viant/ptrie"
)

func main() {
	// 创建新的前缀树
	trie := ptrie.New()
	
	// 添加键值对
	pairs := map[string]interface{}{
		"apple":  1,
		"app":    2,
		"banana": 3,
	}
	
	for key, value := range pairs {
		if err := trie.Put([]byte(key), value); err != nil {
			log.Fatal(err)
		}
	}
	
	// 检查键是否存在
	has := trie.Has([]byte("apple"))
	fmt.Printf("Has 'apple': %v\n", has) // true
	
	// 获取键对应的值
	value, has := trie.Get([]byte("app"))
	fmt.Printf("Get 'app': %v, exists: %v\n", value, has) // 2, true
	
	// 匹配所有前缀
	input := []byte("application")
	matched := trie.MatchAll(input, func(key []byte, value interface{}) bool {
		fmt.Printf("matched: key: %s, value %v\n", key, value)
		return true
	})
	fmt.Printf("Total matches: %d\n", matched)
}

2. 构建和序列化

package main

import (
	"bytes"
	"log"
	"github.com/viant/ptrie"
)

func main() {
	// 创建并填充前缀树
	trie := ptrie.New()
	pairs := map[string]interface{}{
		"go":     "language",
		"golang": "awesome",
		"python": "scripting",
	}
	
	for key, value := range pairs {
		if err := trie.Put([]byte(key), value); err != nil {
			log.Fatal(err)
		}
	}
	
	// 序列化前缀树
	writer := new(bytes.Buffer)
	if err := trie.Encode(writer); err != nil {
		log.Fatal(err)
	}
	encoded := writer.Bytes()
	
	// encoded 可以保存到文件或云存储
}

3. 加载和反序列化

package main

import (
	"bytes"
	"log"
	"reflect"
	"github.com/viant/ptrie"
)

func main() {
	// 假设我们有一个序列化的前缀树
	var encoded []byte // 从文件或存储加载
	
	// 创建新的前缀树并指定值类型
	trie := ptrie.New()
	var v *string // 指定值的类型
	trie.UseType(reflect.TypeOf(v))
	
	// 反序列化
	reader := bytes.NewReader(encoded)
	if err := trie.Decode(reader); err != nil {
		log.Fatal(err)
	}
	
	// 现在可以使用反序列化的前缀树
	value, has := trie.Get([]byte("golang"))
	log.Printf("Get 'golang': %v, exists: %v", value, has)
}

4. 遍历前缀树

package main

import (
	"fmt"
	"github.com/viant/ptrie"
)

func main() {
	trie := ptrie.New()
	pairs := map[string]interface{}{
		"cat":  1,
		"car":  2,
		"cart": 3,
	}
	
	for key, value := range pairs {
		_ = trie.Put([]byte(key), value)
	}
	
	// 遍历所有键值对
	trie.Walk(func(key []byte, value interface{}) bool {
		fmt.Printf("key: %s, value %v\n", key, value)
		return true // 返回true继续遍历，false停止
	})
}

5. 前缀匹配

package main

import (
	"fmt"
	"github.com/viant/ptrie"
)

func main() {
	trie := ptrie.New()
	pairs := map[string]interface{}{
		"app":      1,
		"apple":    2,
		"applet":   3,
		"banana":   4,
		"band":     5,
	}
	
	for key, value := range pairs {
		_ = trie.Put([]byte(key), value)
	}
	
	input := []byte("apple pie")
	
	// 匹配输入的所有前缀
	matched := trie.MatchPrefix(input, func(key []byte, value interface{}) bool {
		fmt.Printf("Prefix matched: key: %s, value %v\n", key, value)
		return true
	})
	fmt.Printf("Total prefix matches: %d\n", matched)
}

性能基准测试

基准测试统计了以下文本中的所有单词：

Lorem Ipsum

1. 短文本：平均行大小20，单词13
2. 长文本：平均行大小711，单词551

Benchmark_LoremBruteForceShort-8    	  500000	      3646 ns/op
Benchmark_LoremTrieShort-8          	  500000	      2376 ns/op
Benchmark_LoremBruteForceLong-8     	    1000	   1612877 ns/op
Benchmark_LoremTrieLong-8           	   10000	    119990 ns/op

Hamlet

1. 短文本：平均行大小20，单词49
2. 长文本：平均行大小41，单词105

Benchmark_HamletBruteForceShort-8   	   30000	     44306 ns/op
Benchmark_HamletTrieShort-8         	  100000	     18530 ns/op
Benchmark_HamletBruteForceLong-8    	   10000	    226836 ns/op
Benchmark_HamletTrieLong-8          	   50000	     39329 ns/op

许可证

源代码根据Apache License, Version 2的条款提供，如LICENSE文件中所述。

作者

库作者: Adrian Witas