Search code examples
rustvector

Accesing elements from Vec using Vec of indices

I'm coming from python to Rust and I want to access very specific elements of given vector using another vector containing indices. So that's python code that I have in mind:

import numpy as np
a = np.array([1,2,3,4,5,6,7])
b = np.array([1,0,2,1])
print(a[b])

resulting in [2 1 3 2] output.

So far I'm able to obtain not so elegant solution to my problem in Rust:

fn main() {
    let a: Vec<i32> = vec![1,2,3,4,5,6,7];
    let b = vec![a[1],a[0],a[2],a[1]]; 
    println!("{:?}", b);
}

resulting in [2 1 3 2] output.

I could also achieve it by writing a for loop, however I hope there's a nicer solution to my problem.

Solution

  • You can use an iterator using map to do this:

    fn main() {
    let a = vec![1, 2, 3, 4, 5, 6, 7];
    let b = vec![1, 0, 2, 1, 8];
    let r = b.iter().map(|&i| a.get(i).map(|v| *v)).collect::<Vec<_>>();
    println!("{:?}", r);
}

    Result:

    [Some(2), Some(1), Some(3), Some(2), None]