Skip to content

Commit 42cac22

Browse files
New Itertools::tail
Note that all cases consume the iterator (even when n=0). It could alternatively return a vector and not `VecIntoIter` but most of our methods do this (and immediately collect to a vector won't reallocate). I don't think we should write a non-lazy `head` method (`.take(n)` is a lazy one) but it would be easy: `.take(n).collect_vec().into_iter()`.
1 parent 8ed734b commit 42cac22

File tree

2 files changed

+52
-0
lines changed

2 files changed

+52
-0
lines changed

src/lib.rs

+47
Original file line numberDiff line numberDiff line change
@@ -3133,6 +3133,53 @@ pub trait Itertools: Iterator {
31333133
self.k_largest_by(k, k_smallest::key_to_cmp(key))
31343134
}
31353135

3136+
/// Consumes the iterator and return an iterator of the last `n` elements.
3137+
///
3138+
/// It allocates up to `n` elements.
3139+
/// The iterator, if directly collected to a `Vec`, is converted
3140+
/// without any extra copying or allocation cost.
3141+
///
3142+
/// ```
3143+
/// use itertools::{assert_equal, Itertools};
3144+
///
3145+
/// let v = vec![5, 9, 8, 4, 2, 12, 0];
3146+
/// assert_equal(v.iter().tail(3), &[2, 12, 0]);
3147+
/// assert_equal(v.iter().tail(10), &v);
3148+
///
3149+
/// assert_equal((0..100).tail(10), 90..100);
3150+
/// ```
3151+
///
3152+
/// For double ended iterators without side-effects, you might prefer
3153+
/// `.rev().take(n).collect_vec().into_iter().rev()`
3154+
/// to have the same result without consuming the entire iterator.
3155+
#[cfg(feature = "use_alloc")]
3156+
fn tail(mut self, n: usize) -> VecIntoIter<Self::Item>
3157+
where
3158+
Self: Sized,
3159+
{
3160+
match n {
3161+
0 => {
3162+
self.last();
3163+
Vec::new()
3164+
}
3165+
1 => self.last().into_iter().collect(),
3166+
_ => {
3167+
let mut data: Vec<_> = self.by_ref().take(n).collect();
3168+
// If the iterator is not exhausted yet, update the data cyclically and
3169+
// finally rotate the data so that the order of the tail is respected.
3170+
if data.len() == n {
3171+
let idx = self.fold(0, |i, val| {
3172+
data[i] = val;
3173+
(i + 1) % n
3174+
});
3175+
data.rotate_left(idx);
3176+
}
3177+
data
3178+
}
3179+
}
3180+
.into_iter()
3181+
}
3182+
31363183
/// Collect all iterator elements into one of two
31373184
/// partitions. Unlike [`Iterator::partition`], each partition may
31383185
/// have a distinct type.

tests/quick.rs

+5
Original file line numberDiff line numberDiff line change
@@ -1949,4 +1949,9 @@ quickcheck! {
19491949
result_set.is_empty()
19501950
}
19511951
}
1952+
1953+
fn tail(v: Vec<i32>, n: u8) -> bool {
1954+
let n = n as usize;
1955+
itertools::equal(v.iter().tail(n), &v[v.len().saturating_sub(n)..])
1956+
}
19521957
}

0 commit comments

Comments
 (0)