Say I've got a Vec that looks something like [[1,2,3],[4,5,6],[7,8,9]]. How would I iterate over all the inner vecs at once? That is,

for k in some_func([[1,2,3],[4,5,6],[7,8,9]]) {
    println!("{k:?}")
}
=> [1,4,7],[2,5,8],[3,6,9]

I can manually construct the transpose, but that only works if the inner iterators are finite. How would I do this for any iterator?

Edit: more scientifically, I want a version of

fn parallel_iterate<T>(items: Vec<Vec<T>>) -> Vec<Vec<T>>
where
    T: Clone,
{
    let size = items.iter().map(|row| row.len()).min().unwrap_or(0);
    let mut result = Vec::with_capacity(size);
    for index in 0..size {
        result.push(Vec::new());
        for inner in items.iter() {
            result[index].push(inner[index].clone())
        }
    }
    return result;
}

where items is of type Vec<I> and I: IntoIterator<Item = T>

Say I've got a Vec that looks something like [[1,2,3],[4,5,6],[7,8,9]]. How would I iterate over all the inner vecs at once? That is,

for k in some_func([[1,2,3],[4,5,6],[7,8,9]]) {
    println!("{k:?}")
}
=> [1,4,7],[2,5,8],[3,6,9]

I can manually construct the transpose, but that only works if the inner iterators are finite. How would I do this for any iterator?

Edit: more scientifically, I want a version of

fn parallel_iterate<T>(items: Vec<Vec<T>>) -> Vec<Vec<T>>
where
    T: Clone,
{
    let size = items.iter().map(|row| row.len()).min().unwrap_or(0);
    let mut result = Vec::with_capacity(size);
    for index in 0..size {
        result.push(Vec::new());
        for inner in items.iter() {
            result[index].push(inner[index].clone())
        }
    }
    return result;
}

where items is of type Vec<I> and I: IntoIterator<Item = T>

• rust
• vector
• iteration

• What do you expect to be returned by the iterator? – fdan Commented Feb 5 at 21:08

4 Answers 4

After the question has been edited, it seems that the number of rows and columns could be arbitrary.
The previous attempt (kept under the below line) considered only a fixed number of rows.

The example provided in the question makes extensive use of vectors, but we can generalise to any kind of iterator.
The main advantage is that we do not need to allocate the transposed result as a whole; we can just consume the partial results on the fly if the whole storage is not necessary.

The idea here is to build our own TransposeIterator structure, keeping the provided row-iterators, in order to consume all of them at each invocation of .next().
The use of IntoIterator instead of Iterator is just a convenience, so that we can directly pass structured data instead of explicitly making iterators at the call site.

mod transpose {
    pub struct TransposeIterator<I: IntoIterator> {
        iters: Vec<I::IntoIter>,
    }
    impl<I: IntoIterator> TransposeIterator<I> {
        pub fn new(iters: impl IntoIterator<Item = I>) -> Self {
            let iters = iters.into_iter().map(|i| i.into_iter());
            Self {
                iters: Vec::from_iter(iters),
            }
        }
    }
    impl<I: IntoIterator> Iterator for TransposeIterator<I> {
        type Item = Vec<I::Item>;

        fn next(&mut self) -> Option<Self::Item> {
            let mut result = Vec::with_capacity(self.iters.len());
            for it in self.iters.iter_mut() {
                if let Some(elem) = it.next() {
                    result.push(elem);
                } else {
                    return None;
                }
            }
            Some(result)
        }
    }
}

fn main() {
    use transpose::TransposeIterator;
    let mut data =
        vec![[11, 12, 13], [21, 22, 23], [31, 32, 33], [41, 42, 43]];
    for row in data.iter() {
        println!("{:?}", row);
    }
    println!("~~~~ consumes original data ~~~~");
    for row in TransposeIterator::new(data.clone()) {
        println!("{:?}", row);
    }
    println!("~~~~ references original data ~~~~");
    for row in TransposeIterator::new(&data) {
        println!("{:?}", row);
    }
    println!("~~~~ mutably references original data ~~~~");
    for mut row in TransposeIterator::new(&mut data) {
        for elem in row.iter_mut() {
            **elem += 900;
        }
        println!("{:?}", row);
    }
}
/*
[11, 12, 13]
[21, 22, 23]
[31, 32, 33]
[41, 42, 43]
~~~~ consumes original data ~~~~
[11, 21, 31, 41]
[12, 22, 32, 42]
[13, 23, 33, 43]
~~~~ references original data ~~~~
[11, 21, 31, 41]
[12, 22, 32, 42]
[13, 23, 33, 43]
~~~~ mutably references original data ~~~~
[911, 921, 931, 941]
[912, 922, 932, 942]
[913, 923, 933, 943]
*/

Initial answer

I would do it with two nested invocations of std::iter::zip().

The best solution is certainly itertools::multizip() (it yields tuples, not arrays, I guess this is not important).

fn some_func<I: IntoIterator>(
    i0: I,
    i1: I,
    i2: I,
) -> impl Iterator<Item = [I::Item; 3]> {
    std::iter::zip(i0, std::iter::zip(i1, i2))
        .map(|(e1, (e2, e3))| [e1, e2, e3])
}

fn main() {
    for k in std::iter::zip([1, 2, 3], std::iter::zip([4, 5, 6], [7, 8, 9])) {
        println!("{k:?}")
    }
    println!("~~~~ consumes original data ~~~~~~~~~~~~");
    for k in some_func([1, 2, 3], [4, 5, 6], [7, 8, 9]) {
        println!("{k:?}")
    }
    println!("~~~~ references original data ~~~~~~~~~~~~");
    for k in some_func(&[1, 2], &[4, 5], &[7, 8]) {
        println!("{k:?}")
    }
    println!("~~~~ consumes original data ~~~~~~~~~~~~");
    for k in itertools::multizip(([1, 2, 3], [4, 5, 6], [7, 8, 9])) {
        println!("{k:?}")
    }
    println!("~~~~ references original data ~~~~~~~~~~~~");
    for k in itertools::multizip((&[1, 2], &[4, 5], &[7, 8])) {
        println!("{k:?}")
    }
}
/*
(1, (4, 7))
(2, (5, 8))
(3, (6, 9))
~~~~ consumes original data ~~~~~~~~~~~~
[1, 4, 7]
[2, 5, 8]
[3, 6, 9]
~~~~ references original data ~~~~~~~~~~~~
[1, 4, 7]
[2, 5, 8]
~~~~ consumes original data ~~~~~~~~~~~~
(1, 4, 7)
(2, 5, 8)
(3, 6, 9)
~~~~ references original data ~~~~~~~~~~~~
(1, 4, 7)
(2, 5, 8)
*/

This satisfies the signature of the edited question:

pub fn parallel_iterate<T>(into_iters: Vec<impl IntoIterator<Item = T>>) -> Vec<Vec<T>> {
    let mut iters: Vec<_> = into_iters
        .into_iter()
        .map(IntoIterator::into_iter)
        .collect();
    let mut result = vec![];
    'outer: loop {
        let mut inner_result = Vec::with_capacity(iters.len());
        for inner_iter in &mut iters {
            let Some(item) = inner_iter.next() else {
                break 'outer;
            };
            inner_result.push(item);
        }
        result.push(inner_result);
    }
    result
}

Playground

Something like this should work.

fn zip_iters<T>(
    v: &mut [impl Iterator<Item = T>],
) -> impl '_ + Iterator<Item = Vec<T>> {
    std::iter::repeat(()).map_while(|()| {
        v.iter_mut().map(|it| it.next()).collect::<Option<Vec<_>>>()
    })
}

fn main() {
    let mut v = [
        // squares of even numbers
        Box::new((0_i32..).filter_map(|i| (i % 2 == 0).then_some(i.pow(2))))
            as Box<dyn Iterator<Item = i32>>,
        // 0..=4
        Box::new((0_i32..).take_while(|&i| i < 5)) as _,
    ];

    println!("{:?}", zip_iters(&mut v).collect::<Vec<_>>());
    // [[0, 0], [4, 1], [16, 2], [36, 3], [64, 4]]
}

How about this?

use core::panic;

fn flatten_vectors<T: Copy>(vectors: Vec<Vec<T>>) -> Vec<Vec<T>> {
    let mut result: Vec<Vec<T>> = Vec::new();
    let max_len = vectors.iter().map(|v| v.len()).max().unwrap_or(0);

    for i in 0..max_len {
        let mut inner = Vec::<T>::new();
        for vec in &vectors {
            if let Some(&val) = vec.get(i) {
                inner.push(val);
            }
            else {
                panic!("Index out of bounds");
            }
        }
        result.push(inner);
    }

    result
}

fn main() {
    let a = [[1, 2, 3], [4, 5, 6], [7, 8, 9]];
    let v: Vec<Vec<i32>> = a.iter().map(|a| a.to_vec()).collect();
    let flattened = flatten_vectors(v);
    println!("{:?}", flattened);
}

-> [[1, 4, 7], [2, 5, 8], [3, 6, 9]]

Playground