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
16pub fn remove_indices<T>(vector: &mut Vec<T>, indices: &[usize]) {
23 debug_assert!(indices.len() <= vector.len());
24 debug_assert!(indices.is_sorted());
25 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() { i += 1;
36 false
37 } else { true
39 };
40
41 vector_index += 1;
42
43 keep
44 });
45}
46
47pub fn remove_sparse_indices<T>(vector: &mut Vec<(usize, T)>, indices: &[usize]) {
56 debug_assert!(indices.is_sorted());
57 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
79pub 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
145pub 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 let mut tuples: Vec<(usize, i32)> = vec![];
267 remove_sparse_indices(&mut tuples, &[]);
268 assert_eq!(tuples, vec![]);
269
270 let mut tuples = vec![(1_000, 1f64)];
272 remove_sparse_indices(&mut tuples, &[]);
273 assert_eq!(tuples, vec![(1_000, 1f64)]);
274
275 let mut tuples: Vec<(usize, i32)> = vec![];
277 remove_sparse_indices(&mut tuples, &[0, 1, 1_000]);
278 assert_eq!(tuples, vec![]);
279
280 let mut tuples = vec![(0, 3f64)];
282 remove_sparse_indices(&mut tuples, &[1]);
283 assert_eq!(tuples, vec![(0, 3f64)]);
284
285 let mut tuples = vec![(2, 3f64)];
287 remove_sparse_indices(&mut tuples, &[1]);
288 assert_eq!(tuples, vec![(1, 3f64)]);
289
290 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 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 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 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 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 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}