使用Trie数据结构高效搜索固定长度字节数组的前缀

package main

import "fmt"

// Fixed 定义了固定长度的字节数组，例如64字节
type Fixed [64]byte

// TrieNode 代表Trie树中的一个节点
type TrieNode struct {
    children map[byte]*TrieNode // 子节点映射，键为字节，值为子节点指针
    values   []Fixed            // 存储以当前节点为完整路径终点的Fixed数组
}

// NewTrieNode 创建一个新的Trie节点
func NewTrieNode() *TrieNode {
    return &TrieNode{
        children: make(map[byte]*TrieNode),
        values:   make([]Fixed, 0),
    }
}

// Trie 代表前缀树
type Trie struct {
    root *TrieNode // Trie的根节点
}

// NewTrie 创建一个新的Trie
func NewTrie() *Trie {
    return &Trie{
        root: NewTrieNode(),
    }
}

// Insert 将一个Fixed数组插入到Trie中
func (t *Trie) Insert(data Fixed) {
    node := t.root
    for i := 0; i < len(data); i++ { // 遍历Fixed数组的每一个字节
        b := data[i]
        if _, ok := node.children[b]; !ok {
            node.children[b] = NewTrieNode() // 如果子节点不存在，则创建
        }
        node = node.children[b] // 移动到下一个节点
    }
    node.values = append(node.values, data) // 将完整的Fixed数组存储在终端节点
}

// FindPrefix 查找所有以给定前缀开头的Fixed数组
func (t *Trie) FindPrefix(prefix []byte) []Fixed {
    node := t.root
    for _, b := range prefix { // 遍历前缀的每一个字节
        if _, ok := node.children[b]; !ok {
            return nil // 如果前缀路径中断，则无匹配项
        }
        node = node.children[b] // 移动到下一个节点
    }

    // 'node' 现在是所有匹配该前缀的Fixed数组的根节点
    var results []Fixed
    t.collectAllValues(node, &results) // 收集该子树中的所有Fixed数组
    return results
}

// collectAllValues 递归地收集从给定节点开始的所有子树中的Fixed数组
func (t *Trie) collectAllValues(node *TrieNode, results *[]Fixed) {
    *results = append(*results, node.values...) // 添加当前节点存储的Fixed数组
    for _, child := range node.children {
        t.collectAllValues(child, results) // 递归收集子节点中的Fixed数组
    }
}

func main() {
    myTrie := NewTrie()

    // 插入一些示例数据
    data1 := Fixed{1, 2, 3, 4, 5, 6, 7, 8, 0, 0 /*... rest of 64 bytes*/}
    data2 := Fixed{1, 2, 3, 4, 5, 6, 7, 9, 0, 0 /*...*/}
    data3 := Fixed{1, 2, 3, 4, 5, 8, 0, 0, 0, 0 /*...*/}
    data4 := Fixed{1, 2, 3, 4, 6, 0, 0, 0, 0, 0 /*...*/}
    data5 := Fixed{10, 11, 12, 0, 0, 0, 0, 0, 0, 0 /*...*/}

    myTrie.Insert(data1)
    myTrie.Insert(data2)
    myTrie.Insert(data3)
    myTrie.Insert(data4)
    myTrie.Insert(data5)

    // 进行前缀搜索
    prefix1 := []byte{1, 2, 3, 4, 5, 6, 7}
    fmt.Printf("Searching for prefix %v:\n", prefix1)
    results1 := myTrie.FindPrefix(prefix1)
    for _, item := range results1 {
        fmt.Printf("  Found: %v\n", item[:8]) // 打印前8个字节作为示例
    }
    // Expected: data1, data2

    prefix2 := []byte{1, 2, 3, 4, 5}
    fmt.Printf("\nSearching for prefix %v:\n", prefix2)
    results2 := myTrie.FindPrefix(prefix2)
    for _, item := range results2 {
        fmt.Printf("  Found: %v\n", item[:8])
    }
    // Expected: data1, data2, data3

    prefix3 := []byte{10, 11}
    fmt.Printf("\nSearching for prefix %v:\n", prefix3)
    results3 := myTrie.FindPrefix(prefix3)
    for _, item := range results3 {
        fmt.Printf("  Found: %v\n", item[:8])
    }
    // Expected: data5

    prefix4 := []byte{99} // 不存在的
    fmt.Printf("\nSearching for prefix %v:\n", prefix4)
    results4 := myTrie.FindPrefix(prefix4)
    if results4 == nil {
        fmt.Println("  No items found.")
    }
    // Expected: No items found.
}