如何移除Rust BTreeMap中已找到的元素？性能与选型疑问

示例代码（整理后） #

use std::collections::BTreeMap;

struct BankAccount {
    balance: i32
}

fn charge_fee(
    active_accounts: &mut BTreeMap<u32, BankAccount>, acct_id: u32, fee: i32) {
    if let Some(tmp) = active_accounts.get_mut(&acct_id) {
        (*tmp).balance -= fee;
        if (*tmp).balance < 0 {
            active_accounts.remove(&acct_id);
        }
    }
}

fn test_charge_fee(account_id: u32, fees: i32) {
    let mut active_accounts = BTreeMap::new();
    active_accounts.insert(0, BankAccount{balance: 100});
    active_accounts.insert(1, BankAccount{balance: 50});
    active_accounts.insert(2, BankAccount{balance: 29});
    active_accounts.insert(3, BankAccount{balance: 87});

    println!("Active accts before charging fees:");
    for (k, v) in active_accounts.iter() {
        println!("{}: {}", k, v.balance);
    }

    charge_fee(&mut active_accounts, account_id, fees);

    println!("Active accts after charging fees:");
    for (k, v) in active_accounts.iter() {
        println!("{}: {}", k, v.balance);
    }
}

问题1：是否存在无需重复查找即可移除键值对的更优方式？ #

有两种高效方案，可将查找次数从两次降低到一次：

方案1：使用take取出元素后处理 #

通过take一次性完成查找并取出元素，修改余额后判断是否重新插入：

fn charge_fee(active_accounts: &mut BTreeMap<u32, BankAccount>, acct_id: u32, fee: i32) {
    if let Some(mut account) = active_accounts.take(&acct_id) {
        account.balance -= fee;
        if account.balance >= 0 {
            active_accounts.insert(acct_id, account);
        }
    }
}

方案2：使用Entry API #

利用Entry枚举的Occupied变体，直接对已存在的条目完成修改和移除操作：

fn charge_fee(active_accounts: &mut BTreeMap<u32, BankAccount>, acct_id: u32, fee: i32) {
    if let std::collections::btree_map::Entry::Occupied(mut occupied_entry) = active_accounts.entry(acct_id) {
        let balance = &mut occupied_entry.get_mut().balance;
        *balance -= fee;
        if *balance < 0 {
            occupied_entry.remove();
        }
    }
}

问题2：BTreeMap是否为此类场景的最佳容器？有无替代方案建议？ #

BTreeMap的适用性取决于业务需求：

• 如果需要键的有序性（如按账户ID排序遍历、范围查询）或稳定内存布局，BTreeMap是合适的选择，操作复杂度为O(log n)。
• 如果仅需要快速单键查找/插入/删除，不需要有序性，HashMap更优——平均操作复杂度为O(1)，性能高于BTreeMap。
• 如果账户ID是连续整数且范围可控，可使用Vec<Option<BankAccount>>，通过索引直接访问，操作复杂度为O(1)，性能最优，但无法灵活处理非连续ID场景。

问题3：我的示例代码是否存在错误？ #

代码本身可正常运行，但有两点需要注意：

1. 安全性无问题：get_mut返回的引用在remove调用时已超出生命周期（if let块结束后tmp被销毁），不会出现悬垂引用或并发修改错误。
2. 功能补充：
   • test_charge_fee函数未被调用，需在main中添加调用逻辑才能运行，例如fn main() { test_charge_fee(2, 30); }。
   • 逻辑符合“保留余额≥0账户”的要求，若业务需移除余额为0的账户，需将判断条件修改为(*tmp).balance <= 0。

内容的提问来源于stack exchange，提问作者jj-loves-programming