Rust中多BTreeSet并集操作的高效实现方法问询

你的实现方式确实会因为频繁创建中间BTreeSet带来额外的内存分配和元素拷贝开销，尤其是当集合规模较大时，效率损耗会更明显。下面提供两种更高效的实现思路，核心都是利用惰性求值避免中间集合的创建：

方法一：借助itertools合并有序迭代器（推荐） #

由于BTreeMap的键是有序的，其keys()返回的迭代器也是有序的。我们可以直接合并这些有序迭代器，再去重后一次性收集为最终的BTreeSet，全程只有一次内存分配。

首先在Cargo.toml中添加依赖：

[dependencies]
itertools = "0.12"

示例代码：

use itertools::Itertools;
use std::collections::{BTreeMap, BTreeSet};

fn main() {
    let btree_map_1: BTreeMap<i32, String> = [(1, "a".into()), (3, "c".into())].into();
    let btree_map_2: BTreeMap<i32, String> = [(2, "b".into()), (3, "c".into())].into();
    let btree_map_3: BTreeMap<i32, String> = [(4, "d".into()), (1, "a".into())].into();

    // 收集所有BTreeMap的键迭代器
    let key_iters = [
        btree_map_1.keys(),
        btree_map_2.keys(),
        btree_map_3.keys(),
    ];

    // 惰性合并有序迭代器、去重，最后一次性收集为BTreeSet
    let merged_set: BTreeSet<_> = key_iters
        .into_iter()
        .flatten()
        .merge() // 合并多个有序迭代器，保持元素有序
        .dedup() // 跳过连续重复的元素
        .cloned() // 克隆引用为所有权类型
        .collect();

    println!("{:?}", merged_set); // 输出: {1, 2, 3, 4}
}

优势 #

• merge和dedup都是惰性操作，不会提前生成中间集合，仅在迭代时处理元素。
• 最终只做一次collect，仅分配一次内存存储最终集合，大幅降低开销。

方法二：手动实现多路归并去重（无第三方依赖） #

如果无法引入第三方库，可以用BinaryHeap实现多路归并，手动合并有序迭代器并去重，同样避免中间集合的创建。

示例代码：

use std::collections::{BTreeMap, BTreeSet, BinaryHeap};
use std::cmp::Reverse;

fn main() {
    let btree_map_1: BTreeMap<i32, String> = [(1, "a".into()), (3, "c".into())].into();
    let btree_map_2: BTreeMap<i32, String> = [(2, "b".into()), (3, "c".into())].into();
    let btree_map_3: BTreeMap<i32, String> = [(4, "d".into()), (1, "a".into())].into();

    // 将迭代器包装为Reverse，把BinaryHeap转为最小堆（默认是最大堆）
    let mut heap: BinaryHeap<_> = [
        Reverse(btree_map_1.keys().peekable()),
        Reverse(btree_map_2.keys().peekable()),
        Reverse(btree_map_3.keys().peekable()),
    ]
    .into_iter()
    .filter(|rev| rev.0.peek().is_some()) // 过滤空迭代器
    .collect();

    let mut merged_set = BTreeSet::new();
    let mut last_val: Option<&i32> = None;

    while let Some(mut rev_iter) = heap.pop() {
        if let Some(&current) = rev_iter.0.peek() {
            // 仅当当前元素与上一个不同时插入，避免重复
            if last_val != Some(current) {
                merged_set.insert(*current);
                last_val = Some(current);
            }
            // 推进迭代器
            rev_iter.0.next();
            // 若迭代器还有元素，放回堆中继续参与合并
            if rev_iter.0.peek().is_some() {
                heap.push(rev_iter);
            }
        }
    }

    println!("{:?}", merged_set); // 输出: {1, 2, 3, 4}
}

优势 #

• 不依赖任何第三方库，纯标准库实现。
• 同样通过惰性迭代处理元素，仅在必要时插入最终集合，无中间集合开销。

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