Hyper-optimize IO in rust solution

1 files changed, 22 insertions, 12 deletions

diff --git a/kattis-kth-alginda-quicksort/src/radix_sort.rs b/kattis-kth-alginda-quicksort/src/radix_sort.rs
index e7d4e02..2f8d32a 100644
--- a/kattis-kth-alginda-quicksort/src/radix_sort.rs
+++ b/kattis-kth-alginda-quicksort/src/radix_sort.rs
@@ -35,18 +35,19 @@ impl RadixFilters for u64 {
     ];
 }
 
-pub fn radix_sort<T>(xs: Vec<T>) -> Vec<T>
+pub fn radix_sort_with<T>(xs: &mut [T], temp: &mut [T])
 where
     T: Default + Copy + Clone + RadixFilters,
 {
     // The first thing that happens it that these two are swapped so thats why this seems backwards
     let mut output = xs;
-    let mut input = vec![T::default(); output.len()];
+    let mut input = temp;
+    debug_assert!(T::FILTERS.len() % 2 == 0);
     for f in T::FILTERS {
         mem::swap(&mut input, &mut output);
 
         let mut counts = [0; 256];
-        for x in &input {
+        for x in input.iter() {
             counts[f(x) as usize] += 1;
         }
         let mut sum = 0;
@@ -60,7 +61,14 @@ where
             output[counts[fx]] = *x;
         }
     }
-    output
+}
+
+pub fn radix_sort<T>(xs: &mut [T])
+where
+    T: Default + Copy + Clone + RadixFilters,
+{
+    let n = xs.len();
+    radix_sort_with(xs, vec![Default::default(); n].as_mut_slice());
 }
 
 #[cfg(test)]
@@ -70,9 +78,9 @@ mod tests {
     #[test]
     fn basic_rand() {
         let mut r = rand_pcg::Pcg32::new(0xcafef00dd15ea5e5, 0xa02bdbf7bb3c0a7);
-        let xs: Vec<u32> = (0..1_000_000).map(|_| r.next_u32()).collect();
+        let mut xs: Vec<u32> = (0..1_000_000).map(|_| r.next_u32()).collect();
 
-        let xs = radix_sort(xs);
+        radix_sort(&mut xs);
         for w in xs.as_slice().windows(2) {
             assert!(w[0] <= w[1]);
         }
@@ -81,9 +89,9 @@ mod tests {
     #[test]
     fn negatives() {
         let mut r = rand_pcg::Pcg32::new(0xcafef00dd15ea5e5, 0xa02bdbf7bb3c0a7);
-        let xs: Vec<i32> = (0..1_000_000).map(|_| r.next_u32() as i32).collect();
+        let mut xs: Vec<i32> = (0..1_000_000).map(|_| r.next_u32() as i32).collect();
 
-        let xs = radix_sort(xs);
+        radix_sort(&mut xs);
         for w in xs.as_slice().windows(2) {
             assert!(w[0] <= w[1], "{} <= {}", w[0], w[1]);
         }
@@ -111,8 +119,9 @@ mod tests {
         let xs: Vec<u32> = (0..1_000_000).map(|_| r.next_u32()).collect();
 
         b.iter(|| {
-            let xs = black_box(xs.clone());
-            black_box(radix_sort(xs));
+            let mut xs = black_box(xs.clone());
+            radix_sort(&mut xs);
+            black_box(xs);
         });
     }
 
@@ -134,8 +143,9 @@ mod tests {
         let xs: Vec<u64> = (0..1_000_000).map(|_| r.next_u64()).collect();
 
         b.iter(|| {
-            let xs = black_box(xs.clone());
-            black_box(radix_sort(xs));
+            let mut xs = black_box(xs.clone());
+            radix_sort(&mut xs);
+            black_box(xs);
         });
     }
 }