Nth with multiple sorted indexes?

1

For some iterators I'd like a function like nth() that accepts an ordered sequence of more than one index. This is useful to avoid scanning the same iterable more than once.

A eager and rough version of it:

fn multi_nth<It1, It2>(mut seq: It1, idxs: It2) -> Vec<u64>
where It1: Iterator<Item=u64>, It2: Iterator<Item=usize> {
    let mut result = vec![];
    let mut si = 0;

    'OUTER: for i in idxs {
        while si <= i {
            if let Some(x) = seq.next() {
                if si == i { result.push(x); }
                si += 1;
            } else {
                break 'OUTER;
            }
        }
    }
    result
}

Do you know where to find a lazy version of this, and if it's a good idea to add it to std lib or itertools?

1 Like
2

Here's a rough version I just threw together: playground

pub struct MultiNth<I, N> {
    iter: I,
    indexes: N,
    iter_idx: usize
}

impl <I: Iterator, N:Iterator<Item=usize>> MultiNth<I, N> {
    fn new(iter: I, indexes: N) -> Self {
        MultiNth {iter, indexes, iter_idx: 0}
    }
}

impl<I: Iterator, N:Iterator<Item=usize>> Iterator for MultiNth<I, N> {
    type Item = I::Item;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        match self.indexes.next() {
            None => None,
            Some(idx) => {
                let mut res = match self.iter.next() {
                    None => return None,
                    Some(next) => next
                };
                self.iter_idx += 1;
                while self.iter_idx < idx {
                    self.iter_idx += 1;
                    res = match self.iter.next() {
                        None => return None,
                        Some(next) => next
                    };
                }
                Some(res)
            }
        }
    }
}
1 Like
3

I think it would be good for itertools

4

I've tried to improve your code a little, now it's usable in an iterator chain. But the handling of indexes is still wrong, and now I'm too much sleepy to debug it (debugging from other people is welcome):

#[derive(Debug)]
struct MultiNthState<I, N>
    where I: Iterator
{
    indexes: N,
    iter_idx: usize,
    underlying: I,
}

trait MultiNth: Iterator {
    fn multi_nth<N>(self, indexes: N) -> MultiNthState<Self, N>
        where N: Iterator<Item=usize>,
        Self: Sized,
    {
        MultiNthState { indexes, iter_idx: 0, underlying: self }
    }
}

impl<I> MultiNth for I where I: Iterator {}

impl<I, N> Iterator for MultiNthState<I, N>
    where I: Iterator,
          N: Iterator<Item=usize>,
{
    type Item = I::Item;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        match self.indexes.next() {
            None => None,
            Some(idx) => {
                let mut res = match self.underlying.next() {
                    None => return None,
                    Some(next) => next
                };
                self.iter_idx += 1;
                while self.iter_idx < idx {
                    self.iter_idx += 1;
                    res = match self.underlying.next() {
                        None => return None,
                        Some(next) => next
                    };
                }
                Some(res)
            }
        }
    }
}

fn main() {
    // Indexes must be ordered and strictly increasing?
    for f in [10, 20, 30, 40, 50, 60, 70]
             .iter()
             .multi_nth([1, 2, 5].iter().cloned()) {
        println!("{}", f); // Should print: 20 30 60
    }
}
5

Alright, I figured out how to actually make it work by just using enumerate to do the index tracking for us:
playground

use std::iter::Enumerate;

pub struct MultiNth<I, N> {
    iter: I,
    indexes: N,
}

impl <I: Iterator, N:Iterator<Item=usize>> MultiNth<Enumerate<I>, N> {
    fn new(iter: I, indexes: N) -> Self {
        MultiNth {iter: iter.enumerate(), indexes}
    }
}

impl<B, I: Iterator<Item=(usize, B)>, N:Iterator<Item=usize>> Iterator for MultiNth<I, N> {
    type Item = B;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        match self.indexes.next() {
            None => None,
            Some(idx) => {
                loop {
                    match self.iter.next() {
                        None => return None,
                        Some((idx2, next)) => {
                            if (idx2 == idx) {
                                return Some(next)
                            }
                        }
                    }
                }
            }
        }
    }
}
6

Repeated application of .nth() could help some iterators turn this into a faster operation.

But then a good way to handle errors is needed (indices not sorted)?

7

A possible solution is to assume the indexes are few (and they are Copy), you can define an Increasing struct with a smart constructor that gives the error:

[10, 20, 30, 40, 50, 60, 70]
.iter()
.multi_nth(Increasing([1, 2, 5].iter().cloned()).unwrap())

Other ideas are welcome.

8

Here's a simple implementation using nth: playground.

pub struct MultiNth<I, N> {
    iter: I,
    indexes: N,
    offset: usize,
}

impl <I: Iterator, N:Iterator<Item=usize>> MultiNth<I, N> {
    fn new(iter: I, indexes: N) -> Self {
        MultiNth {iter, indexes, offset: 0}
    }
}

impl<I: Iterator, N:Iterator<Item=usize>> Iterator for MultiNth<I, N> {
    type Item = I::Item;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        match self.indexes.next() {
            None => None,
            Some(idx) => {
                let o = self.offset;
                let res = self.iter.nth(idx - o);
                self.offset = idx + 1;
                res
            }
        }
    }
}

You could keep a prev field in the Multi struct to detect if an out of date index is given.

1 Like
9

But the point of multi_nth is to avoid calling nth many times on the input sequence.

10

Well, it keeps the Iterator and offset, so it's not scanning from the beginning every time. So complexity-wise it should have the same efficiency.

1 Like
11

Using nth makes it (algorithmically at least) faster with iterators that can turn that into an O(1) step instead of O(n). For example slice iterators.

4 Likes