使用Trie实现固定长度字节数组的高效前缀搜索

package main

import (
    "fmt"
)

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

// TrieNode 代表Trie树中的一个节点
type TrieNode struct {
    Children map[byte]*TrieNode // 子节点，键为字节，值为子节点指针
    IsEndOfWord bool           // 标记是否是某个Fixed数组的结束
    Values []Fixed             // 存储以当前路径为前缀的完整Fixed数组
}

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

// Trie 结构体包含根节点
type Trie struct {
    Root *TrieNode
}

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

// Insert 方法将一个Fixed数组插入到Trie中
func (t *Trie) Insert(data Fixed) {
    node := t.Root
    for _, b := range data {
        if _, ok := node.Children[b]; !ok {
            node.Children[b] = NewTrieNode()
        }
        node = node.Children[b]
    }
    node.IsEndOfWord = true
    node.Values = append(node.Values, data) // 将完整数据存储在结束节点
}

// findNode 方法查找给定前缀对应的节点
func (t *Trie) findNode(prefix []byte) *TrieNode {
    node := t.Root
    for _, b := range prefix {
        if _, ok := node.Children[b]; !ok {
            return nil // 未找到前缀
        }
        node = node.Children[b]
    }
    return node
}

// collectAllValues 从指定节点开始，递归收集所有子树中的Fixed数组
func (t *Trie) collectAllValues(node *TrieNode, results *[]Fixed) {
    if node == nil {
        return
    }
    if node.IsEndOfWord {
        *results = append(*results, node.Values...)
    }
    for _, child := range node.Children {
        t.collectAllValues(child, results)
    }
}

// FindPrefix 方法根据给定的前缀查找所有匹配的Fixed数组
func (t *Trie) FindPrefix(prefix []byte) []Fixed {
    node := t.findNode(prefix)
    if node == nil {
        return nil // 没有匹配的前缀
    }

    var results []Fixed
    // 从前缀节点开始，收集所有以该前缀开头的Fixed数组
    t.collectAllValues(node, &results)
    return results
}

func main() {
    myTrie := NewTrie()

    // 示例数据
    data1 := Fixed{0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09 /*..., other 55 bytes */}
    data2 := Fixed{0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x10, 0x11 /*..., other 55 bytes */}
    data3 := Fixed{0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x08, 0x20, 0x21 /*..., other 55 bytes */}
    data4 := Fixed{0x01, 0x02, 0x03, 0x04, 0x05, 0x09, 0x0A, 0x30, 0x31 /*..., other 55 bytes */}
    data5 := Fixed{0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09 /*..., other 55 bytes */} // Duplicate of data1

    // 填充Trie
    myTrie.Insert(data1)
    myTrie.Insert(data2)
    myTrie.Insert(data3)
    myTrie.Insert(data4)
    myTrie.Insert(data5) // 插入重复数据

    // 查找前缀
    prefix1 := []byte{0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07} // 匹配data1, data2
    results1 := myTrie.FindPrefix(prefix1)
    fmt.Printf("查找前缀 %x 的结果 (%d 个):\n", prefix1, len(results1))
    for _, res := range results1 {
        fmt.Printf("  %x...\n", res[:9])
    }
    // 预期输出：data1, data2

    prefix2 := []byte{0x01, 0x02, 0x03, 0x04, 0x05, 0x06} // 匹配data1, data2, data3
    results2 := myTrie.FindPrefix(prefix2)
    fmt.Printf("\n查找前缀 %x 的结果 (%d 个):\n", prefix2, len(results2))
    for _, res := range results2 {
        fmt.Printf("  %x...\n", res[:9])
    }
    // 预期输出：data1, data2, data3

    prefix3 := []byte{0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x09} // 无匹配
    results3 := myTrie.FindPrefix(prefix3)
    fmt.Printf("\n查找前缀 %x 的结果 (%d 个):\n", prefix3, len(results3))
    if len(results3) == 0 {
        fmt.Println("  无匹配项")
    }

    prefix4 := []byte{0x01, 0x02, 0x03, 0x04, 0x05, 0x09} // 匹配data4
    results4 := myTrie.FindPrefix(prefix4)
    fmt.Printf("\n查找前缀 %x 的结果 (%d 个):\n", prefix4, len(results4))
    for _, res := range results4 {
        fmt.Printf("  %x...\n", res[:9])
    }
}