How to efficiently apply an arbitrary function on a moving window in a BTreeMap-CodePudding

How can I efficiently apply a function over a moving window of a BTreeMap, where the window is determined by a range of the key?

My current code is like this but it is horribly slow when the BTreeMap gets very large (and probably with increasing radius). I suspect it's because of the clone() but I can't think of an alternative solution. How do I do this efficiently?

use nalgebra::Vector3;
use std::collections::BTreeMap;

fn get_window(data: &BTreeMap<i64, Vector3<f64>>, center: &i64, radius: &i64) -> BTreeMap<i64, Vector3<f64>> {
    let mut window = data.clone();
    window.retain(|&ts, _| ts >= (center - radius) && ts <= (center   radius));
    window
}

fn kernel(data: BTreeMap<i64, Vector3<f64>>) -> Vector3<f64> {
    // This is just an example function
    let n = data.len() as f64;
    let x = data.iter().map(|(_, vec)| vec[0]).sum::<f64>() / n;
    let y = data.iter().map(|(_, vec)| vec[1]).sum::<f64>() / n;
    let z = data.iter().map(|(_, vec)| vec[2]).sum::<f64>() / n;
    Vector3::new(x, y, z)
}

fn main() {
    let mut btreemap = BTreeMap::<i64, Vector3<f64>>::new();
    btreemap.insert(0, Vector3::new(0.0, 0.0, 0.0));
    btreemap.insert(1, Vector3::new(1.0, 0.0, 0.0));
    btreemap.insert(2, Vector3::new(0.0, 1.0, 0.0));
    btreemap.insert(3, Vector3::new(0.0, 0.0, 1.0));
    btreemap.insert(4, Vector3::new(2.0, 0.0, 0.0));
    btreemap.insert(5, Vector3::new(0.0, 2.0, 0.0));
    btreemap.insert(6, Vector3::new(0.0, 0.0, 2.0));
    
    let mut smoothed = BTreeMap::<i64, Vector3<f64>>::new();
    for (ts, _vec) in btreemap.iter() {
        let window = get_window(&btreemap, &ts, &1);
        let kernel = kernel(window);
        smoothed.insert(*ts, kernel);
    }

    println!("{:?}", smoothed);
}

CodePudding user response：

Instead of cloning the map and then filtering the items in the window, you can use BTreeMap::range to get an iterator over the items within a range of the map. This is cheap because a btree_map::Range iterates over the data in the source map; it doesn't clone any of it.

I recommend also passing the center and radius as i64 rather than &i64. There's no need to borrow them.

use std::collections::BTreeMap;
use std::collections::btree_map;

fn get_window(data: &BTreeMap<i64, Vector3<f64>>, center: i64, radius: i64) -> btree_map::Range<'_, i64, Vector3<f64>> {
    data.range(center-radius..=center radius)
}

fn kernel(data: btree_map::Range<'_, i64, Vector3<f64>>) -> Vector3<f64> {
    // This is just an example function
    let n = data.clone().count() as f64;
    let x = data.clone().map(|(_, vec)| vec[0]).sum::<f64>() / n;
    let y = data.clone().map(|(_, vec)| vec[1]).sum::<f64>() / n;
    let z = data.clone().map(|(_, vec)| vec[2]).sum::<f64>() / n;
    Vector3::new(x, y, z)
}

The loop in main becomes:

for (&ts, _vec) in btreemap.iter() {
    let window = get_window(&btreemap, ts, 1);
    let kernel = kernel(window);
    smoothed.insert(ts, kernel);
}

Playground