Remove a range from Rust BTreeMap

Viewed 233

I have a BTreeMap<f64,_> and want to leave numbers x where low < x < high. However split_off(k) cannot control k being inclusive or exclusive.

This is a temporary solution with eps:

// delete all numbers >= high
let _ = map.split_off(high);

// want to delete all numbers <= low
map = map.split_off(low+1e-8);

A correct approach is to delete elements one by one, which I think can take O(n lg n). Is there a way to use map.range((Excluded(&low), Excluded(&high))) and delete the range in O(lg n)? It looks like I need map.erase(range) or map.split_off_after(k).

1 Answers

Rust 1.53 introduced retain method for BTreeMap.

pub fn retain<F>(&mut self, f: F) where
    K: Ord,
    F: FnMut(&K, &mut V) -> bool, 

Thus, we can write:

use std::collections::BTreeMap;

fn main() {
    let mut map = BTreeMap::<i64, f64>::new();  // Note, f64 hasn't implemented Ord trait.

    for x in 0..10 {
        map.insert(x, x as f64 * 100.0);
    }

    map.retain(|k, _| *k > 3 && *k < 6);  // <--- remove a range with prediction

    for (k, v) in map {
        println!("{} {}", k, v); // 4 400\n 5 500
    }
}
Related