light_curve_common/
sort_multiple.rs

1use std::cmp::PartialOrd;
2
3/// Sorts multiple slices in order of the first one
4///
5/// # Examples
6///
7/// ```
8/// use light_curve_common::sort_multiple;
9///
10/// let a = [1, 2, 3, 0, 4];
11/// let b = [3, 2, 1, 4, 0];
12///
13/// let sorted = sort_multiple(&[&a, &b]);
14/// let a_sorted = &sorted[0];
15/// let b_sorted = &sorted[1];
16///
17/// assert_eq!(&vec![0, 1, 2, 3, 4], a_sorted);
18/// assert_eq!(&vec![4, 3, 2, 1, 0], b_sorted);
19/// ```
20pub fn sort_multiple<T: PartialOrd + Copy>(slices: &[&[T]]) -> Vec<Vec<T>> {
21    if slices.is_empty() {
22        return vec![];
23    }
24
25    let reference = slices[0];
26    let other = &slices[1..];
27
28    let n = reference.len();
29    for s in other {
30        assert_eq!(s.len(), n);
31    }
32
33    let mut order: Vec<_> = (0..n).collect();
34    order.sort_unstable_by(|&a, &b| reference[a].partial_cmp(&reference[b]).unwrap());
35
36    slices.iter().map(|s| reorder(s, &order)).collect()
37}
38
39fn reorder<T: Copy>(s: &[T], order: &[usize]) -> Vec<T> {
40    order.iter().map(|&i| s[i]).collect()
41}
42
43#[cfg(test)]
44mod tests {
45    use super::*;
46
47    use rand::prelude::*;
48
49    fn sort_via_zip_unzip<T: PartialOrd + Clone>(a: &[T], b: &[T]) -> Vec<Vec<T>> {
50        let mut pairs: Vec<_> = a.iter().cloned().zip(b.iter().cloned()).collect();
51        pairs.sort_unstable_by(|x, y| x.0.partial_cmp(&y.0).unwrap());
52        let (sorted_a, sorted_b) = pairs.into_iter().unzip();
53        vec![sorted_a, sorted_b]
54    }
55
56    #[test]
57    fn two_random() {
58        const N: usize = 100;
59
60        let mut rng = rand::thread_rng();
61        let a: Vec<f64> = (0..N).map(|_| rng.gen()).collect();
62        let b: Vec<f64> = (0..N).map(|_| rng.gen()).collect();
63
64        assert_eq!(sort_via_zip_unzip(&a, &b), sort_multiple(&[&a, &b]));
65    }
66
67    #[test]
68    fn three() {
69        let a = [5, 4, 3, 2, 1, 0];
70        let b = [5, 4, 3, 2, 1, 0];
71        let c = [0, 1, 2, 3, 4, 5];
72
73        assert_eq!(
74            vec![
75                vec![0, 1, 2, 3, 4, 5],
76                vec![0, 1, 2, 3, 4, 5],
77                vec![5, 4, 3, 2, 1, 0]
78            ],
79            sort_multiple(&[&a, &b, &c])
80        );
81    }
82}
light_curve_common/sort_multiple.rs

light_curve_common/
sort_multiple.rs
