Get the number of elements stored inside an n-dimensional vector

Question

I have a 2-dimensional vector:

let vec2d = vec![
    vec![1, 1, 1],
    vec![1, 1, 1],
];

I can yield the total of elements stored this way:

let mut n_vec_element: i32 = 0;

for i in vec2d.iter() {
    n_vec_element += i.len() as i32;
}

println!("2D vector elements :{}", n_vec_element); // prints 6

When I increase the dimensions, the loop gets longer:

let mut n_vec_element: i32 = 0;

let vec3d = vec![
    vec![
        vec![1, 3, 5 as i32],
        vec![2, 4, 6 as i32],
        vec![3, 5, 7 as i32],
    ],
    vec![
        vec![1, 3, 5 as i32],
        vec![2, 4, 6 as i32],
        vec![3, 5, 7 as i32],
    ]
];

for i in vec3d.iter() {

    // I must add another iter everytime I increment the dimension by 1.
    // Else, it returns the number of stored vector instead of the vector 
    // elements.

    for j in i.iter() { 
        n_vec_size += j.len() as i32;
    }
};

println!("3D vector elements :{}", n_vec_element); // prints 18

There must be a more concise way to do this, but I haven't figured it out yet. Initially, I tried using vector's len() function but as I said above, it returns number of vectors stored instead of its elements.

Answer 1

You do not need an explicit loop to do so:

let vec2d = vec![
    vec![1, 1, 1],
    vec![1, 1, 1],
];

let n_vec_element: usize = vec2d.iter().map(Vec::len).sum();

assert_eq!(n_vec_element, 6);

For a 3d vector, you can do the same:

let vec3d = vec![
    vec![
        vec![1, 3, 5 as i32],
        vec![2, 4, 6 as i32],
        vec![3, 5, 7 as i32],
    ],
    vec![
        vec![1, 3, 5 as i32],
        vec![2, 4, 6 as i32],
        vec![3, 5, 7 as i32],
    ]
];

let n_vec_element: usize = vec3d.iter().flatten().map(Vec::len).sum();

assert_eq!(n_vec_element, 18);

With a 4D vector, you can put 2 flatten, etc.

---

With the specialization feature (i.e., with the nightly compiler), you can generalize this with an unique method:

#![feature(specialization)]

trait Count {
    fn count(self) -> usize;
}

impl<T> Count for T {
    default fn count(self) -> usize {
        1
    }
}

impl<T> Count for T
where
    T: IntoIterator,
    T::Item: Count,
{
    fn count(self) -> usize {
        self.into_iter().map(|x| x.count()).sum()
    }
}

fn main() {
    let v = vec![1, 2, 3];
    assert_eq!(v.count(), 3);

    let v = vec![vec![1, 2, 3], vec![4, 5, 6]];
    assert_eq!(v.count(), 6);
}