arr_rs/core/operations/
count.rs

1use crate::{
2    core::prelude::*,
3    errors::prelude::*,
4    extensions::prelude::*,
5};
6
7/// `ArrayTrait` - Array Count functions
8pub trait ArrayCount<T: ArrayElement> where Self: Sized + Clone {
9
10    /// Sort an array
11    ///
12    /// # Arguments
13    ///
14    /// * `axis` - axis along which to count. if None, array is flattened
15    /// * `keepdims` - if true, the result will broadcast correctly against the input
16    ///
17    /// # Examples
18    ///
19    /// ```
20    /// use arr_rs::prelude::*;
21    ///
22    /// let arr = array_eye!(i32, 4);
23    /// assert_eq!(array!(usize, [4]), arr.count_nonzero(None, None));
24    ///
25    /// let arr = array!(i32, [[0, 1, 7, 0], [3, 0, 2, 19]]);
26    /// assert_eq!(array!(usize, [1, 1, 2, 1]), arr.count_nonzero(Some(0), None));
27    /// ```
28    ///
29    /// # Errors
30    ///
31    /// may returns `ArrayError`
32    fn count_nonzero(&self, axis: Option<isize>, keepdims: Option<bool>) -> Result<Array<usize>, ArrayError>;
33}
34
35impl <T: ArrayElement> ArrayCount<T> for Array<T> {
36
37    fn count_nonzero(&self, axis: Option<isize>, keepdims: Option<bool>) -> Result<Array<usize>, ArrayError> {
38        if let Some(axis) = axis {
39            let axis = self.normalize_axis(axis);
40            let result = self.apply_along_axis(axis, |arr| arr.count_nonzero(None, keepdims));
41            if keepdims == Some(true) { result }
42            else { result.reshape(&self.get_shape()?.remove_at(axis)) }
43        } else {
44            let result = Array::single(self
45                .get_elements()?
46                .into_iter()
47                .filter(|e| e != &T::zero()).count());
48            if keepdims == Some(true) { result.atleast(self.ndim()?) }
49            else { result }
50        }
51    }
52}
53
54impl <T: ArrayElement> ArrayCount<T> for Result<Array<T>, ArrayError> {
55
56    fn count_nonzero(&self, axis: Option<isize>, keepdims: Option<bool>) -> Result<Array<usize>, ArrayError> {
57        self.clone()?.count_nonzero(axis, keepdims)
58    }
59}