smartcore/algorithm/sort/
quick_sort.rs

1use num_traits::Num;
2
3pub trait QuickArgSort {
4    fn quick_argsort_mut(&mut self) -> Vec<usize>;
5
6    #[allow(dead_code)]
7    fn quick_argsort(&self) -> Vec<usize>;
8}
9
10impl<T: Num + PartialOrd + Copy> QuickArgSort for Vec<T> {
11    fn quick_argsort(&self) -> Vec<usize> {
12        let mut v = self.clone();
13        v.quick_argsort_mut()
14    }
15
16    fn quick_argsort_mut(&mut self) -> Vec<usize> {
17        let stack_size = 64;
18        let mut jstack = -1;
19        let mut l = 0;
20        let mut istack = vec![0; stack_size];
21        let mut ir = self.len() - 1;
22        let mut index: Vec<usize> = (0..self.len()).collect();
23
24        loop {
25            if ir - l < 7 {
26                for j in l + 1..=ir {
27                    let a = self[j];
28                    let b = index[j];
29                    let mut i: i32 = (j - 1) as i32;
30                    while i >= l as i32 {
31                        if self[i as usize] <= a {
32                            break;
33                        }
34                        self[(i + 1) as usize] = self[i as usize];
35                        index[(i + 1) as usize] = index[i as usize];
36                        i -= 1;
37                    }
38                    self[(i + 1) as usize] = a;
39                    index[(i + 1) as usize] = b;
40                }
41                if jstack < 0 {
42                    break;
43                }
44                ir = istack[jstack as usize];
45                jstack -= 1;
46                l = istack[jstack as usize];
47                jstack -= 1;
48            } else {
49                let k = (l + ir) >> 1;
50                self.swap(k, l + 1);
51                index.swap(k, l + 1);
52                if self[l] > self[ir] {
53                    self.swap(l, ir);
54                    index.swap(l, ir);
55                }
56                if self[l + 1] > self[ir] {
57                    self.swap(l + 1, ir);
58                    index.swap(l + 1, ir);
59                }
60                if self[l] > self[l + 1] {
61                    self.swap(l, l + 1);
62                    index.swap(l, l + 1);
63                }
64                let mut i = l + 1;
65                let mut j = ir;
66                let a = self[l + 1];
67                let b = index[l + 1];
68                loop {
69                    loop {
70                        i += 1;
71                        if self[i] >= a {
72                            break;
73                        }
74                    }
75                    loop {
76                        j -= 1;
77                        if self[j] <= a {
78                            break;
79                        }
80                    }
81                    if j < i {
82                        break;
83                    }
84                    self.swap(i, j);
85                    index.swap(i, j);
86                }
87                self[l + 1] = self[j];
88                self[j] = a;
89                index[l + 1] = index[j];
90                index[j] = b;
91                jstack += 2;
92
93                if jstack >= 64 {
94                    panic!("stack size is too small.");
95                }
96
97                if ir - i + 1 >= j - l {
98                    istack[jstack as usize] = ir;
99                    istack[jstack as usize - 1] = i;
100                    ir = j - 1;
101                } else {
102                    istack[jstack as usize] = j - 1;
103                    istack[jstack as usize - 1] = l;
104                    l = i;
105                }
106            }
107        }
108
109        index
110    }
111}
112
113#[cfg(test)]
114mod tests {
115    use super::*;
116
117    #[cfg_attr(
118        all(target_arch = "wasm32", not(target_os = "wasi")),
119        wasm_bindgen_test::wasm_bindgen_test
120    )]
121    #[test]
122    fn with_capacity() {
123        let arr1 = vec![0.3, 0.1, 0.2, 0.4, 0.9, 0.5, 0.7, 0.6, 0.8];
124        assert_eq!(vec![1, 2, 0, 3, 5, 7, 6, 8, 4], arr1.quick_argsort());
125
126        let arr2 = vec![
127            0.2, 0.2, 0.2, 0.2, 0.2, 0.4, 0.3, 0.2, 0.2, 0.1, 1.4, 1.5, 1.5, 1.3, 1.5, 1.3, 1.6,
128            1.0, 1.3, 1.4,
129        ];
130        assert_eq!(
131            vec![9, 7, 1, 8, 0, 2, 4, 3, 6, 5, 17, 18, 15, 13, 19, 10, 14, 11, 12, 16],
132            arr2.quick_argsort()
133        );
134    }
135}
smartcore/algorithm/sort/quick_sort.rs

smartcore/algorithm/sort/
quick_sort.rs
