Skip to main content

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