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    fn scatter_assign(&mut self, ixs: &<VecKind as ArrayKind>::Index, values: Self) {
133        for (i, x) in ixs.iter().zip(values.iter()) {
134            self[*i] = x.clone();
135        }
136    }
137}
138
139impl Add<&VecArray<usize>> for usize {
140    type Output = VecArray<usize>;
141
142    fn add(self, rhs: &VecArray<usize>) -> Self::Output {
143        VecArray(rhs.iter().map(|x| x + self).collect())
144    }
145}
146
147impl<T: Clone + Add<Output = T>> Add<VecArray<T>> for VecArray<T> {
148    type Output = VecArray<T>;
149
150    fn add(self, rhs: VecArray<T>) -> VecArray<T> {
151        assert_eq!(self.len(), rhs.len());
152        VecArray(
153            self.iter()
154                .zip(rhs.iter())
155                .map(|(x, y)| x.clone() + y.clone())
156                .collect(),
157        )
158    }
159}
160
161impl<T: Clone + Sub<Output = T>> Sub<VecArray<T>> for VecArray<T> {
162    type Output = VecArray<T>;
163
164    fn sub(self, rhs: VecArray<T>) -> VecArray<T> {
165        assert_eq!(self.len(), rhs.len());
166        VecArray(
167            self.iter()
168                .zip(rhs.iter())
169                .map(|(x, y)| x.clone() - y.clone())
170                .collect(),
171        )
172    }
173}
174
175impl<T: Ord + Clone> OrdArray<VecKind, T> for VecArray<T> {
176    /// ```rust
177    /// # use open_hypergraphs::array::*;
178    /// # use open_hypergraphs::array::vec::*;
179    /// let values: VecArray<usize> = VecArray(vec![1, 2, 0, 3]);
180    /// let actual: VecArray<usize> = values.argsort();
181    /// let expected = VecArray::<usize>(vec![2, 0, 1, 3]);
182    /// assert_eq!(actual, expected);
183    ///
184    /// // Check monotonicity
185    /// let monotonic = values.gather(actual.as_slice());
186    /// for i in 0..(monotonic.len()-1) {
187    ///     assert!(monotonic[i] <= monotonic[i+1]);
188    /// }
189    /// ```
190    fn argsort(&self) -> VecArray<usize> {
191        let mut indices = (0..self.len()).collect::<Vec<_>>();
192        indices.sort_by_key(|&i| &self[i]);
193        VecArray(indices)
194    }
195}
196
197impl NaturalArray<VecKind> for VecArray<usize> {
198    fn max(&self) -> Option<usize> {
199        self.iter().max().copied()
200    }
201
202    /// ```rust
203    /// # use open_hypergraphs::array::{*, vec::*};
204    /// let x = VecArray(vec![0, 1, 2, 3, 4, 5]);
205    /// let d = 3;
206    /// let expected_q = VecArray(vec![0, 0, 0, 1, 1, 1]);
207    /// let expected_r = VecArray(vec![0, 1, 2, 0, 1, 2]);
208    /// let (q, r) = x.quot_rem(d);
209    /// assert_eq!(expected_q, q);
210    /// assert_eq!(expected_r, r);
211    /// ```
212    fn quot_rem(&self, d: usize) -> (Self, Self) {
213        assert!(d != 0);
214        let mut q = Vec::with_capacity(self.len());
215        let mut r = Vec::with_capacity(self.len());
216        for x in self.iter() {
217            q.push(x / d);
218            r.push(x % d);
219        }
220        (VecArray(q), VecArray(r))
221    }
222
223    fn mul_constant_add(&self, c: usize, x: &Self) -> Self {
224        assert_eq!(self.len(), x.len());
225        let mut r = Vec::with_capacity(self.len());
226        for (s, x) in self.iter().zip(x.iter()) {
227            r.push(s * c + x)
228        }
229        VecArray(r)
230    }
231
232    /// ```rust
233    /// # use open_hypergraphs::array::{*, vec::*};
234    /// let input = VecArray(vec![1, 2, 3, 4]);
235    /// let expected = VecArray(vec![0, 1, 3, 6, 10]);
236    ///
237    /// assert_eq!(input.cumulative_sum(), expected);
238    /// ```
239    fn cumulative_sum(&self) -> Self {
240        let mut v = Vec::with_capacity(self.len() + 1);
241        let mut a = 0;
242        for x in self.iter() {
243            v.push(a);
244            a += x;
245        }
246        v.push(a); // don't forget the total sum!
247        VecArray(v)
248    }
249
250    fn arange(start: &usize, stop: &usize) -> Self {
251        assert!(stop >= start, "invalid range [{:?}, {:?})", start, stop);
252        let n = stop - start;
253        let mut v = Vec::with_capacity(n);
254        for i in 0..n {
255            v.push(start + i);
256        }
257        VecArray(v)
258    }
259
260    /// ```rust
261    /// # use open_hypergraphs::array::*;
262    /// # use open_hypergraphs::array::vec::*;
263    /// let repeats: VecArray<usize> = VecArray(vec![1, 2, 0, 3]);
264    /// let values: &[usize] = &[5, 6, 7, 8];
265    /// let actual = repeats.repeat(values);
266    /// let expected = VecArray::<usize>(vec![5, 6, 6, 8, 8, 8]);
267    /// assert_eq!(actual, expected);
268    /// ```
269    fn repeat(&self, x: &[usize]) -> VecArray<usize> {
270        assert_eq!(self.len(), x.len());
271        let mut v: Vec<usize> = Vec::new();
272        for (k, xi) in self.iter().zip(x) {
273            v.extend(std::iter::repeat(xi).take(*k))
274        }
275        VecArray(v)
276    }
277
278    fn connected_components(
279        sources: &Self,
280        targets: &Self,
281        n: usize,
282    ) -> (Self, <VecKind as ArrayKind>::I) {
283        let (cc_ix, c) = connected_components(sources, targets, n);
284        (VecArray(cc_ix), c)
285    }
286
287    fn bincount(&self, size: usize) -> VecArray<usize> {
288        let mut counts = vec![0; size];
289        for &idx in self.iter() {
290            counts[idx] += 1;
291        }
292        VecArray(counts)
293    }
294
295    fn zero(&self) -> VecArray<usize> {
296        let mut zero_indices = Vec::with_capacity(self.len());
297        for (i, &val) in self.iter().enumerate() {
298            if val == 0 {
299                zero_indices.push(i);
300            }
301        }
302        VecArray(zero_indices)
303    }
304
305    fn sparse_bincount(&self) -> (VecArray<usize>, VecArray<usize>) {
306        use std::collections::HashMap;
307
308        // Count occurrences using a HashMap
309        let mut counts_map = HashMap::new();
310        for &idx in self.iter() {
311            *counts_map.entry(idx).or_insert(0) += 1;
312        }
313
314        // Extract and sort unique indices
315        let mut unique_indices: Vec<_> = counts_map.keys().cloned().collect();
316        unique_indices.sort_unstable();
317
318        // Gather counts in the same order as unique indices
319        let counts: Vec<_> = unique_indices.iter().map(|&idx| counts_map[&idx]).collect();
320
321        (VecArray(unique_indices), VecArray(counts))
322    }
323
324    fn scatter_sub_assign(&mut self, ixs: &VecArray<usize>, rhs: &VecArray<usize>) {
325        for i in 0..ixs.len() {
326            self[ixs[i]] -= rhs[i];
327        }
328    }
329}