CODEWITHSUNDEEP
rust
Cornelius Roemer
Cornelius Roemer

Reputation: 8396

Iterator that skips every nth element

Rather than taking every Nth element from an iterator which I can do with Iterator::step_by, I would like to skip every Nth element. How can I achieve this idiomatically? Is there maybe even a standard library or itertools function?

This is what I came up with to skip every 7th say. It requires enumerate, filter, and map, though one could use a filter_map instead of the latter two.

(0..100).enumerate()
      .filter(|&(i, x)| (i + 1) % 7 != 0)
      .map(|(i, x)| x);

How could I cast this into a function so that I could simply write:

(0..100).skip_every(7)

Upvotes: 0

Views: 575

Answers (2)

Sven Marnach
Sven Marnach

Reputation: 602715

If you want to get the exact interface you asked for, your best option at this time is to implement a custom iterator adapter type. Here's a basic version of such a type:

pub struct SkipEvery<I> {
    inner: I,
    every: usize,
    index: usize,
}

impl<I> SkipEvery<I> {
    fn new(inner: I, every: usize) -> Self {
        assert!(every > 1);
        let index = 0;
        Self {
            inner,
            every,
            index,
        }
    }
}

impl<I: Iterator> Iterator for SkipEvery<I> {
    type Item = I::Item;

    fn next(&mut self) -> Option<Self::Item> {
        if self.index == self.every - 1 {
            self.index = 1;
            self.inner.nth(1)
        } else {
            self.index += 1;
            self.inner.next()
        }
    }
}

pub trait IteratorSkipEveryExt: Iterator + Sized {
    fn skip_every(self, every: usize) -> SkipEvery<Self> {
        SkipEvery::new(self, every)
    }
}

impl<I: Iterator + Sized> IteratorSkipEveryExt for I {}

(Playground)

A more complete implementation could also add optimized versions of further Iterator methods, as well as implementations of DoubleEndedIterator and ExactSizeIterator -- see the implementation of StepBy as an example.

Upvotes: 3

Jmb
Jmb

Reputation: 23463

Your code is pretty easy to turn into a function:

fn skip_every<I: Iterator> (iter: I, n: usize) -> impl Iterator<Item = <I as Iterator>::Item> {
    iter.enumerate()
        .filter_map(move |(i, v)| if (i + 1) % n != 0 { Some (v) } else { None })
}

fn main() {
    println!("{:?}", skip_every (0..20, 7).collect::<Vec<_>>());
}

Playground

Or avoiding the expensive modulo:

fn skip_every2<I: Iterator> (iter: I, n: usize) -> impl Iterator<Item = <I as Iterator>::Item> {
    iter.zip ((0..n).rev().cycle()).filter_map (|(v, i)| if i != 0 { Some (v) } else { None })
}

Playground

Upvotes: 2

Related Questions