index_utils/
lib.rs

1#![feature(is_sorted)]
2
3use std::cmp::Ordering;
4use std::collections::HashSet;
5use std::mem::swap;
6
7#[cfg(feature = "num-traits")]
8mod num;
9#[cfg(feature = "num-traits")]
10pub use num::inner_product;
11#[cfg(feature = "num-traits")]
12pub use num::inner_product_slice;
13#[cfg(feature = "num-traits")]
14pub use num::inner_product_slice_iter;
15
16/// Reduce the size of the vector by removing values.
17///
18/// # Arguments
19///
20/// * `vector`: `Vec` to remove indices from.
21/// * `indices`: A sorted slice of indices to remove.
22pub fn remove_indices<T>(vector: &mut Vec<T>, indices: &[usize]) {
23    debug_assert!(indices.len() <= vector.len());
24    debug_assert!(indices.is_sorted());
25    // All values are unique
26    debug_assert!(indices.iter().collect::<HashSet<_>>().len() == indices.len());
27    debug_assert!(indices.iter().all(|&i| i < vector.len()));
28
29    let mut i = 0;
30    let mut vector_index = 0;
31    vector.retain(|_| {
32        let keep = if i < indices.len() && vector_index < indices[i] {
33            true
34        } else if i < indices.len() { // must have j == to_remove[i]
35            i += 1;
36            false
37        } else { // i == to_remove.len()
38            true
39        };
40
41        vector_index += 1;
42
43        keep
44    });
45}
46
47/// Reduce the size of the vector by removing values.
48///
49/// The method operates in place.
50///
51/// # Arguments
52///
53/// * `vector`: `Vec` to remove indices from.
54/// * `indices`: A sorted slice of indices to remove.
55pub fn remove_sparse_indices<T>(vector: &mut Vec<(usize, T)>, indices: &[usize]) {
56    debug_assert!(indices.is_sorted());
57    // All values are unique
58    debug_assert!(indices.iter().collect::<HashSet<_>>().len() == indices.len());
59
60    if indices.is_empty() || vector.is_empty() {
61        return;
62    }
63
64    let mut nr_skipped_before = 0;
65    vector.retain_mut(|(i, _)| {
66        while nr_skipped_before < indices.len() && indices[nr_skipped_before] < *i {
67            nr_skipped_before += 1;
68        }
69
70        if nr_skipped_before < indices.len() && indices[nr_skipped_before] == *i {
71            false
72        } else {
73            *i -= nr_skipped_before;
74            true
75        }
76    });
77}
78
79/// Collect two sparse iterators by merging them.
80///
81/// # Arguments
82///
83/// * `left`: The first iterator, sorted by index.
84/// * `right`: The second iterator, sorted by index.
85/// * `operation`: A binary operation applied when an element is found at some index in both
86/// iterators.
87/// * `operation_left`: A unitary operation applied when an element is found only in the left
88/// iterator.
89/// * `operation_right`: A unitary operation applied when an element is found only in the right
90/// iterator.
91/// * `is_not_default`: A function specifying whether the result if the binary operation is the
92/// default value (that is often zero).
93///
94/// # Return value
95///
96/// A vector with sparse elements by sorted and unique index.
97pub fn merge_sparse_indices<I: Ord, T: Eq, U>(
98    left: impl Iterator<Item=(I, T)>,
99    right: impl Iterator<Item=(I, T)>,
100    operation: impl Fn(T, T) -> U,
101    operation_left: impl Fn(T) -> U,
102    operation_right: impl Fn(T) -> U,
103    is_not_default: impl Fn(&U) -> bool,
104) -> Vec<(I, U)> {
105    let mut results = Vec::with_capacity(left.size_hint().0.max(right.size_hint().0));
106
107    let mut left = left.peekable();
108    let mut right = right.peekable();
109
110    while let (Some((left_index, _)), Some((right_index, _))) = (left.peek(), right.peek()) {
111        let (index, result) = match left_index.cmp(right_index) {
112            Ordering::Less => {
113                let (index, value) = left.next().unwrap();
114
115                (index, operation_left(value))
116            }
117            Ordering::Equal => {
118                let (left_index, left_value) = left.next().unwrap();
119                let (_right_index, right_value) = right.next().unwrap();
120
121                (left_index, operation(left_value, right_value))
122            }
123            Ordering::Greater => {
124                let (index, value) = right.next().unwrap();
125
126                (index, operation_right(value))
127            }
128        };
129
130        if is_not_default(&result) {
131            results.push((index, result));
132        }
133    }
134
135    for (left_index, value) in left {
136        results.push((left_index, operation_left(value)))
137    }
138    for (right_index, value) in right {
139        results.push((right_index, operation_right(value)))
140    }
141
142    results
143}
144
145/// Collect two sparse iterators by merging the indices that they share, dropping the others.
146///
147/// # Arguments
148///
149/// * `left`: The first iterator, sorted by index.
150/// * `right`: The second iterator, sorted by index.
151/// * `operation`: A binary operation applied when an element is found at some index in both
152/// iterators.
153/// * `is_not_default`: A function specifying whether the result if the binary operation is the
154/// default value (that is often zero).
155///
156/// # Return value
157///
158/// A vector with sparse elements by sorted and unique index.
159pub fn merge_sparse_indices_intersect<I: Ord, T, U>(
160    mut left: impl Iterator<Item=(I, T)>,
161    right: impl Iterator<Item=(I, T)>,
162    operation: impl Fn(T, T) -> U,
163    is_not_default: impl Fn(&U) -> bool,
164) -> Vec<(I, U)> {
165    if let Some((mut left_index, mut left_value)) = left.next() {
166        let mut results = Vec::new();
167
168        for (right_index, right_value) in right {
169            match right_index.cmp(&left_index) {
170                Ordering::Less => {}
171                Ordering::Equal => {
172                    if let Some((mut left_next_index, mut new_left_value)) = left.next() {
173                        let (left_old_index, left_old_value) = {
174                            swap(&mut left_index, &mut left_next_index);
175                            swap(&mut left_value, &mut new_left_value);
176
177                            (left_next_index, new_left_value)
178                        };
179
180                        let result = operation(left_old_value, right_value);
181                        if is_not_default(&result) {
182                            results.push((left_old_index, result));
183                        }
184                    } else {
185                        let result = operation(left_value, right_value);
186                        if is_not_default(&result) {
187                            results.push((left_index, result));
188                        }
189                        break;
190                    }
191                }
192                Ordering::Greater => {
193                    'scope: {
194                        while let Some((left_next_index, new_left_value)) = left.next() {
195                            match left_next_index.cmp(&right_index) {
196                                Ordering::Less => {}
197                                Ordering::Equal => {
198                                    let result = operation(new_left_value, right_value);
199                                    if is_not_default(&result) {
200                                        results.push((left_next_index, result));
201                                    }
202                                    if let Some((left_next_index, new_left_value)) = left.next() {
203                                        left_index = left_next_index;
204                                        left_value = new_left_value;
205                                        break 'scope;
206                                    } else {
207                                        return results;
208                                    }
209                                }
210                                Ordering::Greater => {
211                                    left_index = left_next_index;
212                                    left_value = new_left_value;
213                                    break 'scope;
214                                }
215                            }
216                        }
217
218                        return results;
219                    }
220                }
221            }
222        }
223
224        results
225    } else {
226        Vec::with_capacity(0)
227    }
228}
229
230#[cfg(test)]
231mod test {
232    use std::convert::identity;
233    use std::iter::empty;
234    use std::ops::Add;
235
236    use relp_num::NonZero;
237
238    use crate::{merge_sparse_indices, merge_sparse_indices_intersect, remove_indices, remove_sparse_indices};
239
240    #[test]
241    fn test_remove_indices() {
242        let mut v = vec![0f64, 1f64, 2f64];
243        remove_indices(&mut v, &[1]);
244        assert_eq!(v, vec![0f64, 2f64]);
245
246        let mut v = vec![3f64, 0f64, 0f64];
247        remove_indices(&mut v, &[0]);
248        assert_eq!(v, vec![0f64, 0f64]);
249
250        let mut v = vec![0f64, 0f64, 2f64, 3f64, 0f64, 5f64, 0f64, 0f64, 0f64, 9f64];
251        remove_indices(&mut v, &[3, 4, 6]);
252        assert_eq!(v, vec![0f64, 0f64, 2f64, 5f64, 0f64, 0f64, 9f64]);
253
254        let mut v = vec![0f64];
255        remove_indices(&mut v, &[0]);
256        assert_eq!(v, vec![]);
257
258        let mut v: Vec<i32> = vec![];
259        remove_indices(&mut v, &[]);
260        assert!(v.is_empty());
261    }
262
263    #[test]
264    fn test_remove_sparse_indices() {
265        // Empty vec, removing nothing
266        let mut tuples: Vec<(usize, i32)> = vec![];
267        remove_sparse_indices(&mut tuples, &[]);
268        assert_eq!(tuples, vec![]);
269
270        // Non-empty vec, removing nothing
271        let mut tuples = vec![(1_000, 1f64)];
272        remove_sparse_indices(&mut tuples, &[]);
273        assert_eq!(tuples, vec![(1_000, 1f64)]);
274
275        // Empty vec
276        let mut tuples: Vec<(usize, i32)> = vec![];
277        remove_sparse_indices(&mut tuples, &[0, 1, 1_000]);
278        assert_eq!(tuples, vec![]);
279
280        // Removing value not present, index above
281        let mut tuples = vec![(0, 3f64)];
282        remove_sparse_indices(&mut tuples, &[1]);
283        assert_eq!(tuples, vec![(0, 3f64)]);
284
285        // Removing value not present, index below
286        let mut tuples = vec![(2, 3f64)];
287        remove_sparse_indices(&mut tuples, &[1]);
288        assert_eq!(tuples, vec![(1, 3f64)]);
289
290        // Removing present value, index should be adjusted
291        let mut tuples = vec![(0, 0f64), (2, 2f64)];
292        remove_sparse_indices(&mut tuples, &[0]);
293        assert_eq!(tuples, vec![(1, 2f64)]);
294
295        let mut tuples = vec![(0, 0), (2, 2), (4, 4)];
296        remove_sparse_indices(&mut tuples, &[0, 3]);
297        assert_eq!(tuples, vec![(1, 2), (2, 4)]);
298
299        let mut tuples = vec![(0, 0), (2, 2), (4, 4)];
300        remove_sparse_indices(&mut tuples, &[0, 3, 4]);
301        assert_eq!(tuples, vec![(1, 2)]);
302    }
303
304    #[test]
305    fn test_merge_sparse_indices() {
306        // Empty
307        let left: Vec<(i8, i16)> = vec![];
308        let right = vec![];
309
310        let result = merge_sparse_indices(
311            left.into_iter(),
312            right.into_iter(),
313            Add::add,
314            identity,
315            identity,
316            NonZero::is_not_zero,
317        );
318        let expected = vec![];
319        assert_eq!(result, expected);
320
321        // One empty
322        let left: Vec<(i8, i16)> = vec![(2, 1)];
323        let right = vec![];
324
325        let result = merge_sparse_indices(
326            left.into_iter(),
327            right.into_iter(),
328            Add::add,
329            identity,
330            identity,
331            NonZero::is_not_zero,
332        );
333        let expected = vec![(2, 1)];
334        assert_eq!(result, expected);
335
336        // Not related
337        let left = vec![(1, 6)].into_iter();
338        let right = vec![(4, 9)].into_iter();
339
340        let result = merge_sparse_indices(
341            left.into_iter(),
342            right.into_iter(),
343            Add::add,
344            identity,
345            identity,
346            NonZero::is_not_zero,
347        );
348        let expected = vec![(1, 6), (4, 9)];
349        assert_eq!(result, expected);
350
351        // Same index
352        let left = vec![(1, 6)].into_iter();
353        let right = vec![(1, 9)].into_iter();
354
355        let result = merge_sparse_indices(
356            left.into_iter(),
357            right.into_iter(),
358            Add::add,
359            identity,
360            identity,
361            NonZero::is_not_zero,
362        );
363        let expected = vec![(1, 15)];
364        assert_eq!(result, expected);
365
366        // Alternating
367        let left = vec![(1, 6), (3, 4)].into_iter();
368        let right = vec![(2, 9)].into_iter();
369
370        let result = merge_sparse_indices(
371            left.into_iter(),
372            right.into_iter(),
373            Add::add,
374            identity,
375            identity,
376            NonZero::is_not_zero,
377        );
378        let expected = vec![(1, 6), (2, 9), (3, 4)];
379        assert_eq!(result, expected);
380    }
381
382    #[test]
383    fn test_merge_sparse_indices_intersect() {
384        let result: Vec<(usize, i32)> = merge_sparse_indices_intersect(
385            empty(),
386            empty(),
387            |x: i32, y| x * y,
388            |x| *x != 0,
389        );
390        assert_eq!(result, vec![]);
391
392        let result = merge_sparse_indices_intersect(
393            empty(),
394            [(0, 1)].into_iter(),
395            |x, y| x * y,
396            |x| *x != 0,
397        );
398        assert_eq!(result, vec![]);
399
400        let result = merge_sparse_indices_intersect(
401            [(0, 1)].into_iter(),
402            empty(),
403            |x, y| x * y,
404            |x| *x != 0,
405        );
406        assert_eq!(result, vec![]);
407
408        let result = merge_sparse_indices_intersect(
409            [(0, 1), (2, 1), (3, 1)].into_iter(),
410            [(0, 1), (1, 2), (2, 4), (3, 8)].into_iter(),
411            |x, y| x * y,
412            |x| *x != 0,
413        );
414        assert_eq!(result, vec![(0, 1), (2, 4), (3, 8)]);
415
416        let result = merge_sparse_indices_intersect(
417            [(0, 1), (1, 2), (2, 4), (3, 8)].into_iter(),
418            [(0, 1), (2, 1), (3, 1)].into_iter(),
419            |x, y| x * y,
420            |x| *x != 0,
421        );
422        assert_eq!(result, vec![(0, 1), (2, 4), (3, 8)]);
423
424        let result = merge_sparse_indices_intersect(
425            [(0, 1), (1, 2), (3, 8)].into_iter(),
426            [(2, 1), (3, 1)].into_iter(),
427            |x, y| x * y,
428            |x| *x != 0,
429        );
430        assert_eq!(result, vec![(3, 8)]);
431
432        let result = merge_sparse_indices_intersect(
433            [(0, 1), (1, 2), (3, 8)].into_iter(),
434            [(0, 1), (2, 1), (4, 1)].into_iter(),
435            |x, y| x * y,
436            |x| *x != 0,
437        );
438        assert_eq!(result, vec![(0, 1)]);
439
440        let result = merge_sparse_indices_intersect(
441            [(0, 1), (1, 2), (3, 8), (5, 7)].into_iter(),
442            [(0, 1), (2, 1), (4, 1), (5, 7)].into_iter(),
443            |x, y| x * y,
444            |x| *x != 0,
445        );
446        assert_eq!(result, vec![(0, 1), (5, 49)]);
447    }
448}