Implement argmin argmax

src/quantile.rs

@@ -182,6 +182,31 @@ where
     S: Data<Elem = A>,
     D: Dimension,
 {
+    /// Finds the first index of the minimum value of the array.
+    ///
+    /// Returns `None` if any of the pairwise orderings tested by the function
+    /// are undefined. (For example, this occurs if there are any
+    /// floating-point NaN values in the array.)
+    ///
+    /// Returns `None` if the array is empty.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// extern crate ndarray;
+    /// extern crate ndarray_stats;
+    ///
+    /// use ndarray::array;
+    /// use ndarray_stats::QuantileExt;
+    ///
+    /// let a = array![[1., 3., 5.],
+    ///                [2., 0., 6.]];
+    /// assert_eq!(a.argmin(), Some((1, 1)));
+    /// ```
+    fn argmin(&self) -> Option<D::Pattern>
+    where
+        A: PartialOrd;
+
     /// Finds the elementwise minimum of the array.
     ///
     /// Returns `None` if any of the pairwise orderings tested by the function
@@ -203,6 +228,31 @@ where
         A: MaybeNan,
         A::NotNan: Ord;
 
+    /// Finds the first index of the maximum value of the array.
+    ///
+    /// Returns `None` if any of the pairwise orderings tested by the function
+    /// are undefined. (For example, this occurs if there are any
+    /// floating-point NaN values in the array.)
+    ///
+    /// Returns `None` if the array is empty.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// extern crate ndarray;
+    /// extern crate ndarray_stats;
+    ///
+    /// use ndarray::array;
+    /// use ndarray_stats::QuantileExt;
+    ///
+    /// let a = array![[1., 3., 7.],
+    ///                [2., 5., 6.]];
+    /// assert_eq!(a.argmax(), Some((0, 2)));
+    /// ```
+    fn argmax(&self) -> Option<D::Pattern>
+    where
+        A: PartialOrd;
+
     /// Finds the elementwise maximum of the array.
     ///
     /// Returns `None` if any of the pairwise orderings tested by the function
@@ -278,6 +328,23 @@ where
     S: Data<Elem = A>,
     D: Dimension,
 {
+    fn argmin(&self) -> Option<D::Pattern>
+    where
+        A: PartialOrd,
+    {
+        let mut current_min = self.first()?;
+        let mut current_pattern_min = D::zeros(self.ndim()).into_pattern();
+
+        for (pattern, elem) in self.indexed_iter() {
+            if elem.partial_cmp(current_min)? == cmp::Ordering::Less {
+                current_pattern_min = pattern;
+                current_min = elem
+            }
+        }
+
+        Some(current_pattern_min)
+    }
+
     fn min(&self) -> Option<&A>
     where
         A: PartialOrd,
@@ -303,6 +370,23 @@ where
         }))
     }
 
+    fn argmax(&self) -> Option<D::Pattern>
+    where
+        A: PartialOrd,
+    {
+        let mut current_max = self.first()?;
+        let mut current_pattern_max = D::zeros(self.ndim()).into_pattern();
+
+        for (pattern, elem) in self.indexed_iter() {
+            if elem.partial_cmp(current_max)? == cmp::Ordering::Greater {
+                current_pattern_max = pattern;
+                current_max = elem
+            }
+        }
+
+        Some(current_pattern_max)
+    }
+
     fn max(&self) -> Option<&A>
     where
         A: PartialOrd,

tests/quantile.rs

@@ -8,6 +8,24 @@ use ndarray_stats::{
     Quantile1dExt, QuantileExt,
 };
 
+#[test]
+fn test_argmin() {
+    let a = array![[1, 5, 3], [2, 0, 6]];
+    assert_eq!(a.argmin(), Some((1, 1)));
+
+    let a = array![[1., 5., 3.], [2., 0., 6.]];
+    assert_eq!(a.argmin(), Some((1, 1)));
+
+    let a = array![[1., 5., 3.], [2., ::std::f64::NAN, 6.]];
+    assert_eq!(a.argmin(), None);
+
+    let a = array![[1, 0, 3], [2, 0, 6]];
+    assert_eq!(a.argmin(), Some((0, 1)));
+
+    let a: Array2<i32> = array![[], []];
+    assert_eq!(a.argmin(), None);
+}
+
 #[test]
 fn test_min() {
     let a = array![[1, 5, 3], [2, 0, 6]];
@@ -35,6 +53,24 @@ fn test_min_skipnan_all_nan() {
     assert!(a.min_skipnan().is_nan());
 }
 
+#[test]
+fn test_argmax() {
+    let a = array![[1, 5, 3], [2, 0, 6]];
+    assert_eq!(a.argmax(), Some((1, 2)));
+
+    let a = array![[1., 5., 3.], [2., 0., 6.]];
+    assert_eq!(a.argmax(), Some((1, 2)));
+
+    let a = array![[1., 5., 3.], [2., ::std::f64::NAN, 6.]];
+    assert_eq!(a.argmax(), None);
+
+    let a = array![[1, 5, 6], [2, 0, 6]];
+    assert_eq!(a.argmax(), Some((0, 2)));
+
+    let a: Array2<i32> = array![[], []];
+    assert_eq!(a.argmax(), None);
+}
+
 #[test]
 fn test_max() {
     let a = array![[1, 5, 7], [2, 0, 6]];