1use alloc::{vec, vec::Vec};
2use core::ops::{Index, IndexMut};
3use num::{Bounded, Zero};
4
5#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
6#[derive(Debug, Clone)]
7pub struct Array2<T> {
9 nrows: usize,
10 ncols: usize,
11 data: Vec<T>,
12}
13
14impl<T> Array2<T> {
15 #[inline]
17 pub fn new(nrows: usize, ncols: usize, data: Vec<T>) -> Self {
18 assert_eq!(data.len(), ncols * nrows);
19 Array2 { ncols, nrows, data }
20 }
21
22 #[inline]
24 pub fn repeat(nrows: usize, ncols: usize, elt: T) -> Self
25 where
26 T: Copy,
27 {
28 Array2 {
29 nrows,
30 ncols,
31 data: vec![elt; ncols * nrows],
32 }
33 }
34
35 #[inline]
37 pub fn zeros(nrows: usize, ncols: usize) -> Self
38 where
39 T: Copy + Zero,
40 {
41 Self::repeat(nrows, ncols, T::zero())
42 }
43
44 #[inline]
46 pub fn flat_index(&self, i: usize, j: usize) -> usize {
47 i + j * self.nrows
48 }
49
50 #[inline]
52 pub fn data(&self) -> &[T] {
53 &self.data
54 }
55
56 #[inline]
58 pub fn data_mut(&mut self) -> &mut [T] {
59 &mut self.data
60 }
61
62 #[inline]
64 pub fn nrows(&self) -> usize {
65 self.nrows
66 }
67
68 #[inline]
70 pub fn ncols(&self) -> usize {
71 self.ncols
72 }
73
74 #[inline]
76 pub fn min(&self) -> T
77 where
78 T: Copy + Bounded + PartialOrd,
79 {
80 self.data
81 .iter()
82 .fold(Bounded::max_value(), |a, b| if a < *b { a } else { *b })
83 }
84
85 #[inline]
87 pub fn max(&self) -> T
88 where
89 T: Copy + Bounded + PartialOrd,
90 {
91 self.data
92 .iter()
93 .fold(Bounded::min_value(), |a, b| if a > *b { a } else { *b })
94 }
95
96 #[inline]
98 pub fn from_fn(nrows: usize, ncols: usize, f: impl Fn(usize, usize) -> T) -> Self {
99 let mut data = Vec::with_capacity(ncols * nrows);
100
101 for j in 0..ncols {
102 for i in 0..nrows {
103 data.push(f(i, j));
104 }
105 }
106 data.shrink_to_fit();
107 Self { nrows, ncols, data }
108 }
109}
110
111impl<T> Index<(usize, usize)> for Array2<T> {
112 type Output = T;
113
114 #[inline]
115 fn index(&self, index: (usize, usize)) -> &Self::Output {
116 let fid = self.flat_index(index.0, index.1);
117 &self.data[fid]
118 }
119}
120
121impl<T> IndexMut<(usize, usize)> for Array2<T> {
122 #[inline]
123 fn index_mut(&mut self, index: (usize, usize)) -> &mut Self::Output {
124 let fid = self.flat_index(index.0, index.1);
125 &mut self.data[fid]
126 }
127}