open_hypergraphs/array/vec/
vec_array.rs

1//! [`Vec<T>`]-backed arrays
2use super::connected_components::connected_components;
3use crate::array::*;
4use core::ops::{Add, Deref, DerefMut, Index, RangeBounds, Sub};
5
6/// Arrays backed by a [`Vec<T>`].
7#[derive(PartialEq, Eq, Clone, Debug)]
8pub struct VecKind {}
9
10impl ArrayKind for VecKind {
11    type Type<T> = VecArray<T>;
12    type I = usize;
13    type Index = VecArray<usize>;
14
15    // A Slice for Vec is just a rust slice
16    type Slice<'a, T: 'a> = &'a [T];
17}
18
19/// A newtype wrapper for [`Vec<T>`] allowing pointwise arithmetic operations.
20#[derive(PartialEq, Eq, Clone, Debug)]
21pub struct VecArray<T>(pub Vec<T>);
22
23impl AsRef<<VecKind as ArrayKind>::Index> for VecArray<usize> {
24    fn as_ref(&self) -> &<VecKind as ArrayKind>::Index {
25        self
26    }
27}
28
29impl AsMut<<VecKind as ArrayKind>::Index> for VecArray<usize> {
30    fn as_mut(&mut self) -> &mut <VecKind as ArrayKind>::Index {
31        self
32    }
33}
34
35// VecArray is a newtype wrapper, so we can just treat it like a regular old Vec.
36impl<T> Deref for VecArray<T> {
37    type Target = Vec<T>;
38    fn deref(&self) -> &Self::Target {
39        &self.0
40    }
41}
42
43impl<T> DerefMut for VecArray<T> {
44    fn deref_mut(&mut self) -> &mut Self::Target {
45        &mut self.0
46    }
47}
48
49impl<T: Clone + PartialEq> Array<VecKind, T> for VecArray<T> {
50    fn empty() -> Self {
51        VecArray(Vec::default())
52    }
53
54    fn len(&self) -> usize {
55        self.0.len()
56    }
57
58    fn concatenate(&self, other: &Self) -> Self {
59        let mut result: Vec<T> = Vec::with_capacity(self.len() + other.len());
60        result.extend_from_slice(self);
61        result.extend_from_slice(other);
62        VecArray(result)
63    }
64
65    fn fill(x: T, n: usize) -> Self {
66        VecArray(vec![x; n])
67    }
68
69    fn get(&self, i: usize) -> T {
70        self[i].clone()
71    }
72
73    /// Get a contiguous subrange of the array
74    ///
75    /// ```rust
76    /// use open_hypergraphs::array::{*, vec::*};
77    /// let v = VecArray(vec![0, 1, 2, 3, 4]);
78    /// assert_eq!(v.get_range(..), &[0, 1, 2, 3, 4]);
79    /// assert_eq!(v.get_range(..v.len()), &[0, 1, 2, 3, 4]);
80    fn get_range<R: RangeBounds<usize>>(&self, rb: R) -> &[T] {
81        self.index(self.to_range(rb))
82    }
83
84    fn set_range<R: RangeBounds<usize>>(&mut self, rb: R, v: &<VecKind as ArrayKind>::Type<T>) {
85        let r = self.to_range(rb);
86        self[r].clone_from_slice(v)
87    }
88
89    fn gather(&self, idx: &[usize]) -> Self {
90        VecArray(idx.iter().map(|i| self.0[*i].clone()).collect())
91    }
92
93    /// Scatter values over the specified indices `self[idx[i]] = v[i]`.
94    ///
95    /// ```rust
96    /// use open_hypergraphs::array::{*, vec::*};
97    /// let idx = VecArray(vec![2, 1, 0, 2]);
98    /// let v = VecArray(vec![0, 2, 1, 2]);
99    ///
100    /// let expected = VecArray(vec![1, 2, 2]);
101    ///
102    /// let actual = v.scatter(idx.get_range(..), 3);
103    /// assert_eq!(actual, expected);
104    /// ```
105    fn scatter(&self, idx: &[usize], n: usize) -> VecArray<T> {
106        // If self is empty, we return the empty array because there can be no valid indices
107        if self.is_empty() {
108            assert!(idx.is_empty());
109            return VecArray(vec![]);
110        }
111
112        // Otherwise, we fill the result with an arbitrary value ...
113        let mut y = vec![self[0].clone(); n];
114
115        // ... then scatter values of self into result at indexes idx..
116        for (i, x) in self.iter().enumerate() {
117            y[idx[i]] = x.clone();
118        }
119        VecArray(y)
120    }
121
122    fn from_slice(slice: &[T]) -> Self {
123        VecArray(slice.into())
124    }
125
126    fn scatter_assign_constant(&mut self, ixs: &VecArray<usize>, arg: T) {
127        for &idx in ixs.iter() {
128            self[idx] = arg.clone();
129        }
130    }
131}
132
133impl Add<&VecArray<usize>> for usize {
134    type Output = VecArray<usize>;
135
136    fn add(self, rhs: &VecArray<usize>) -> Self::Output {
137        VecArray(rhs.iter().map(|x| x + self).collect())
138    }
139}
140
141impl<T: Clone + Add<Output = T>> Add<VecArray<T>> for VecArray<T> {
142    type Output = VecArray<T>;
143
144    fn add(self, rhs: VecArray<T>) -> VecArray<T> {
145        assert_eq!(self.len(), rhs.len());
146        VecArray(
147            self.iter()
148                .zip(rhs.iter())
149                .map(|(x, y)| x.clone() + y.clone())
150                .collect(),
151        )
152    }
153}
154
155impl<T: Clone + Sub<Output = T>> Sub<VecArray<T>> for VecArray<T> {
156    type Output = VecArray<T>;
157
158    fn sub(self, rhs: VecArray<T>) -> VecArray<T> {
159        assert_eq!(self.len(), rhs.len());
160        VecArray(
161            self.iter()
162                .zip(rhs.iter())
163                .map(|(x, y)| x.clone() - y.clone())
164                .collect(),
165        )
166    }
167}
168
169impl<T: Ord + Clone> OrdArray<VecKind, T> for VecArray<T> {
170    /// ```rust
171    /// # use open_hypergraphs::array::*;
172    /// # use open_hypergraphs::array::vec::*;
173    /// let values: VecArray<usize> = VecArray(vec![1, 2, 0, 3]);
174    /// let actual: VecArray<usize> = values.argsort();
175    /// let expected = VecArray::<usize>(vec![2, 0, 1, 3]);
176    /// assert_eq!(actual, expected);
177    ///
178    /// // Check monotonicity
179    /// let monotonic = values.gather(actual.as_slice());
180    /// for i in 0..(monotonic.len()-1) {
181    ///     assert!(monotonic[i] <= monotonic[i+1]);
182    /// }
183    /// ```
184    fn argsort(&self) -> VecArray<usize> {
185        let mut indices = (0..self.len()).collect::<Vec<_>>();
186        indices.sort_by_key(|&i| &self[i]);
187        VecArray(indices)
188    }
189}
190
191impl NaturalArray<VecKind> for VecArray<usize> {
192    fn max(&self) -> Option<usize> {
193        self.iter().max().copied()
194    }
195
196    /// ```rust
197    /// # use open_hypergraphs::array::{*, vec::*};
198    /// let x = VecArray(vec![0, 1, 2, 3, 4, 5]);
199    /// let d = 3;
200    /// let expected_q = VecArray(vec![0, 0, 0, 1, 1, 1]);
201    /// let expected_r = VecArray(vec![0, 1, 2, 0, 1, 2]);
202    /// let (q, r) = x.quot_rem(d);
203    /// assert_eq!(expected_q, q);
204    /// assert_eq!(expected_r, r);
205    /// ```
206    fn quot_rem(&self, d: usize) -> (Self, Self) {
207        assert!(d != 0);
208        let mut q = Vec::with_capacity(self.len());
209        let mut r = Vec::with_capacity(self.len());
210        for x in self.iter() {
211            q.push(x / d);
212            r.push(x % d);
213        }
214        (VecArray(q), VecArray(r))
215    }
216
217    fn mul_constant_add(&self, c: usize, x: &Self) -> Self {
218        assert_eq!(self.len(), x.len());
219        let mut r = Vec::with_capacity(self.len());
220        for (s, x) in self.iter().zip(x.iter()) {
221            r.push(s * c + x)
222        }
223        VecArray(r)
224    }
225
226    /// ```rust
227    /// # use open_hypergraphs::array::{*, vec::*};
228    /// let input = VecArray(vec![1, 2, 3, 4]);
229    /// let expected = VecArray(vec![0, 1, 3, 6, 10]);
230    ///
231    /// assert_eq!(input.cumulative_sum(), expected);
232    /// ```
233    fn cumulative_sum(&self) -> Self {
234        let mut v = Vec::with_capacity(self.len() + 1);
235        let mut a = 0;
236        for x in self.iter() {
237            v.push(a);
238            a += x;
239        }
240        v.push(a); // don't forget the total sum!
241        VecArray(v)
242    }
243
244    fn arange(start: &usize, stop: &usize) -> Self {
245        assert!(stop >= start, "invalid range [{:?}, {:?})", start, stop);
246        let n = stop - start;
247        let mut v = Vec::with_capacity(n);
248        for i in 0..n {
249            v.push(start + i);
250        }
251        VecArray(v)
252    }
253
254    /// ```rust
255    /// # use open_hypergraphs::array::*;
256    /// # use open_hypergraphs::array::vec::*;
257    /// let repeats: VecArray<usize> = VecArray(vec![1, 2, 0, 3]);
258    /// let values: &[usize] = &[5, 6, 7, 8];
259    /// let actual = repeats.repeat(values);
260    /// let expected = VecArray::<usize>(vec![5, 6, 6, 8, 8, 8]);
261    /// assert_eq!(actual, expected);
262    /// ```
263    fn repeat(&self, x: &[usize]) -> VecArray<usize> {
264        assert_eq!(self.len(), x.len());
265        let mut v: Vec<usize> = Vec::new();
266        for (k, xi) in self.iter().zip(x) {
267            v.extend(std::iter::repeat(xi).take(*k))
268        }
269        VecArray(v)
270    }
271
272    fn connected_components(
273        sources: &Self,
274        targets: &Self,
275        n: usize,
276    ) -> (Self, <VecKind as ArrayKind>::I) {
277        let (cc_ix, c) = connected_components(sources, targets, n);
278        (VecArray(cc_ix), c)
279    }
280
281    fn bincount(&self, size: usize) -> VecArray<usize> {
282        let mut counts = vec![0; size];
283        for &idx in self.iter() {
284            counts[idx] += 1;
285        }
286        VecArray(counts)
287    }
288
289    fn zero(&self) -> VecArray<usize> {
290        let mut zero_indices = Vec::with_capacity(self.len());
291        for (i, &val) in self.iter().enumerate() {
292            if val == 0 {
293                zero_indices.push(i);
294            }
295        }
296        VecArray(zero_indices)
297    }
298
299    fn sparse_bincount(&self) -> (VecArray<usize>, VecArray<usize>) {
300        use std::collections::HashMap;
301
302        // Count occurrences using a HashMap
303        let mut counts_map = HashMap::new();
304        for &idx in self.iter() {
305            *counts_map.entry(idx).or_insert(0) += 1;
306        }
307
308        // Extract and sort unique indices
309        let mut unique_indices: Vec<_> = counts_map.keys().cloned().collect();
310        unique_indices.sort_unstable();
311
312        // Gather counts in the same order as unique indices
313        let counts: Vec<_> = unique_indices.iter().map(|&idx| counts_map[&idx]).collect();
314
315        (VecArray(unique_indices), VecArray(counts))
316    }
317
318    fn scatter_sub_assign(&mut self, ixs: &VecArray<usize>, rhs: &VecArray<usize>) {
319        for i in 0..ixs.len() {
320            self[ixs[i]] -= rhs[i];
321        }
322    }
323}