Rust中BTreeMap前缀匹配：无需获取下一个字符的实现方案

问题描述 #

需要实现一个函数，从BTreeMap中返回所有匹配给定前缀的字符串及其关联值。当前的实现需要构造前缀的“下一个字符”作为范围查询的上限，但char::from_u32操作可能失败，希望找到无需生成下一个字符的替代方案。

解决方案 #

方案1：利用前缀+最小字符构造范围上限 #

由于BTreeMap的键是按字典序排序的，所有以指定前缀开头的键，必然满足：大于等于前缀，且小于前缀+空字符\0（\0是ASCII码最小的字符，任何常规可见字符的字典序都比它大）。这种方式无需计算“下一个字符”，直接构造上限即可。

针对String类型的实现：

use std::collections::BTreeMap;

fn test_string_prefix_range() {
    let mut map = BTreeMap::new();
    map.insert(String::from("aa"), 1);
    map.insert(String::from("aab"), 2);
    map.insert(String::from("aac"), 4);
    map.insert(String::from("ab"), 5);
    map.insert(String::from("ba"), 6);

    let prefix = "aa";
    // 构造范围上限：前缀拼接空字符
    let upper_bound = format!("{}\0", prefix);
    let matches: Vec<(String, i32)> = map.range(prefix..&upper_bound)
        .map(|(s, val)| (s.clone(), *val))
        .collect();
    
    println!("{:?}", matches); // 输出: [("aa", 1), ("aab", 2), ("aac", 4)]
}

针对Vec类型的实现：

use std::collections::BTreeMap;

fn test_vec_prefix_range() {
    let mut map = BTreeMap::new();
    map.insert(Vec::from(['a', 'a']), 1);
    map.insert(Vec::from(['a', 'a', 'b']), 2);
    map.insert(Vec::from(['a', 'a', 'c']), 4);
    map.insert(Vec::from(['a', 'b']), 5);
    map.insert(Vec::from(['b', 'a']), 5);

    let prefix = vec!['a', 'a'];
    // 构造范围上限：前缀添加空字符
    let mut upper_bound = prefix.clone();
    upper_bound.push('\0');
    let matches: Vec<(Vec<char>, i32)> = map.range(prefix..upper_bound)
        .map(|(chars, val)| (chars.clone(), *val))
        .collect();
    
    println!("{:?}", matches); // 输出: [(['a', 'a'], 1), (['a', 'a', 'b'], 2), (['a', 'a', 'c'], 4)]
}

方案2：遍历过滤（简单但低效） #

如果不想构造范围，可以直接遍历整个BTreeMap，过滤出以指定前缀开头的键值对。这种方法实现简单，但时间复杂度为O(n)，数据量大时性能不如范围查询（范围查询为O(log n + k)，k是匹配的数量）。

use std::collections::BTreeMap;

fn test_filter_prefix() {
    let mut map = BTreeMap::new();
    map.insert(String::from("aa"), 1);
    map.insert(String::from("aab"), 2);
    map.insert(String::from("aac"), 4);
    map.insert(String::from("ab"), 5);
    map.insert(String::from("ba"), 6);

    let prefix = "aa";
    let matches: Vec<(String, i32)> = map.iter()
        .filter(|(key, _)| key.starts_with(prefix))
        .map(|(s, val)| (s.clone(), *val))
        .collect();
    
    println!("{:?}", matches);
}

方案3：使用Trie（前缀树）结构 #

如果需要频繁进行前缀查询，Trie（前缀树）是更适合的数据结构，它本身就是为前缀匹配场景设计的，查询效率更高。可以自行实现简单Trie，或使用第三方库。

简单Trie实现示例：

use std::collections::HashMap;

struct TrieNode {
    value: Option<i32>,
    children: HashMap<char, TrieNode>,
}

impl TrieNode {
    fn new() -> Self {
        TrieNode {
            value: None,
            children: HashMap::new(),
        }
    }

    // 插入键值对到Trie中
    fn insert(&mut self, key: &str, value: i32) {
        let mut node = self;
        for c in key.chars() {
            node = node.children.entry(c).or_insert_with(TrieNode::new);
        }
        node.value = Some(value);
    }

    // 获取所有匹配指定前缀的键值对
    fn get_prefix_matches(&self, prefix: &str) -> Vec<(String, i32)> {
        let mut node = self;
        // 先定位到前缀对应的节点
        for c in prefix.chars() {
            match node.children.get(&c) {
                Some(child) => node = child,
                None => return Vec::new(), // 无匹配前缀，直接返回空
            }
        }

        let mut results = Vec::new();
        self.collect_subtree_matches(node, prefix.to_string(), &mut results);
        results
    }

    // 递归收集子树中的所有键值对
    fn collect_subtree_matches(&self, node: &TrieNode, current_key: String, results: &mut Vec<(String, i32)>) {
        // 如果当前节点有值，加入结果集
        if let Some(val) = node.value {
            results.push((current_key.clone(), val));
        }
        // 遍历所有子节点，继续收集
        for (c, child) in &node.children {
            let mut new_key = current_key.clone();
            new_key.push(*c);
            self.collect_subtree_matches(child, new_key, results);
        }
    }
}

// 使用示例
fn test_trie() {
    let mut trie = TrieNode::new();
    trie.insert("aa", 1);
    trie.insert("aab", 2);
    trie.insert("aac", 4);
    trie.insert("ab", 5);
    trie.insert("ba", 6);

    let matches = trie.get_prefix_matches("aa");
    println!("{:?}", matches); // 输出: [("aa", 1), ("aab", 2), ("aac", 4)]
}

内容的提问来源于stack exchange，提问作者Pioneer_11