diff --git a/src/slice.rs b/src/slice.rs index 32e14a8c7..988b21fd7 100644 --- a/src/slice.rs +++ b/src/slice.rs @@ -48,6 +48,48 @@ impl Slice { pub fn step_by(self, step: Ixs) -> Self { Slice { step, ..self } } + + /// Returns the `n`th index in this slice. + /// + /// The `n` argument starts from zero, so `nth(0)` is the first index. If + /// `n` is greater than or equal to the length of the slice, then the + /// return value is `None`. + /// + /// # Examples + /// + /// ``` + /// use ndarray::Slice; + /// + /// // indices 0, 1, 2, 3, 4 + /// assert_eq!(Slice::new(0, Some(5), 1).nth(2), Some(2)); + /// + /// // indices 1, 3 + /// assert_eq!(Slice::new(1, Some(5), 2).nth(1), Some(3)); + /// + /// // indices 1, -1, -3 + /// assert_eq!(Slice::new(1, Some(-4), -2).nth(2), Some(-3)); + /// + /// // indices 1, 3, 5, ... + /// assert_eq!(Slice::new(1, None, 2).nth(2), Some(5)); + /// + /// // indices 3, 3, 3, ... + /// assert_eq!(Slice::new(3, Some(5), 0).nth(1000), Some(3)); + /// + /// // indices 1, 3 + /// assert_eq!(Slice::new(1, Some(5), 2).nth(2), None); + /// ``` + pub fn nth(&self, n: usize) -> Option { + let nth = self.start + self.step * n as isize; + if let Some(end) = self.end { + if self.step == 0 || (self.step > 0 && nth < end) || (self.step < 0 && nth > end) { + Some(nth) + } else { + None + } + } else { + Some(nth) + } + } } impl From> for Slice {