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