use crate::{
core::prelude::*,
errors::prelude::*,
extensions::prelude::*,
};
pub trait ArrayCount<T: ArrayElement> where Self: Sized + Clone {
fn count_nonzero(&self, axis: Option<isize>, keepdims: Option<bool>) -> Result<Array<usize>, ArrayError>;
}
impl <T: ArrayElement> ArrayCount<T> for Array<T> {
fn count_nonzero(&self, axis: Option<isize>, keepdims: Option<bool>) -> Result<Array<usize>, ArrayError> {
if let Some(axis) = axis {
let axis = self.normalize_axis(axis);
let result = self.apply_along_axis(axis, |arr| arr.count_nonzero(None, keepdims));
if keepdims == Some(true) { result }
else { result.reshape(&self.get_shape()?.remove_at(axis)) }
} else {
let result = Array::single(self
.get_elements()?
.into_iter()
.filter(|e| e != &T::zero()).count());
if keepdims == Some(true) { result.atleast(self.ndim()?) }
else { result }
}
}
}
impl <T: ArrayElement> ArrayCount<T> for Result<Array<T>, ArrayError> {
fn count_nonzero(&self, axis: Option<isize>, keepdims: Option<bool>) -> Result<Array<usize>, ArrayError> {
self.clone()?.count_nonzero(axis, keepdims)
}
}