Hyper-optimize IO in rust solution

4 files changed, 83 insertions, 25 deletions

diff --git a/kattis-kth-alginda-quicksort/radix.c b/kattis-kth-alginda-quicksort/radix.c
index b195d57..93bd111 100644
--- a/kattis-kth-alginda-quicksort/radix.c
+++ b/kattis-kth-alginda-quicksort/radix.c
@@ -62,6 +62,7 @@ int main() {
     }
     while (*p == ' ' || *p == '\n') p++;
 
+    // 18 ms
     for (int i = 0; i < n; i++) {
         int x = 0;
         bool neg = false;
@@ -80,8 +81,10 @@ int main() {
         xs[i] = x ^ (1 << 31);
     }
 
+    // 10 ms
     radix_sort(n);
 
+    // 17 ms
     p = &buffer[BUFFER_MAX - 1];
     char *last = p;
     *p-- = '\n';
diff --git a/kattis-kth-alginda-quicksort/src/main.rs b/kattis-kth-alginda-quicksort/src/main.rs
index 67646b3..afb23f7 100644
--- a/kattis-kth-alginda-quicksort/src/main.rs
+++ b/kattis-kth-alginda-quicksort/src/main.rs
@@ -2,27 +2,71 @@
 
 #![feature(test)]
 
-use std::fmt::Write;
-use std::io::{stdin, Read};
+use std::io::{stdin, stdout, Read, Write};
 
 mod radix_sort;
-
 use radix_sort::radix_sort;
 
 fn main() {
     let mut buffer = String::new();
     stdin().lock().read_to_string(&mut buffer).unwrap();
-    let xs: Vec<i32> = buffer
-        .split_ascii_whitespace()
-        .skip(1)
-        .map(|x| x.parse().unwrap())
-        .collect();
+    let mut p = buffer.bytes();
+    let mut n = 0;
+    while let Some(b) = p.next() {
+        if b == b' ' {
+            break;
+        }
+        n *= 10;
+        n += (b - b'0') as usize;
+    }
+    let mut xs = vec![0; n];
+    for i in 0..n {
+        let mut neg = false;
+        let mut x = 0i32;
+        while let Some(b) = p.next() {
+            match b {
+                b' ' | b'\n' => break,
+                b'-' => neg = true,
+                b => {
+                    x *= 10;
+                    x += (b - b'0') as i32;
+                }
+            }
+        }
+        if neg {
+            x = -x;
+        }
+        xs[i] = x;
+    }
 
-    let xs = radix_sort(xs);
+    radix_sort(&mut xs);
 
-    buffer.clear();
-    for x in xs {
-        writeln!(buffer, "{}", x).unwrap();
+    let mut p = buffer.len() - 1;
+    let b = unsafe { buffer.as_mut_vec() };
+    b[p] = b'\n';
+    p -= 1;
+    for i in (0..n).rev() {
+        let x = xs[i];
+        let neg = x < 0;
+        let mut x = if neg { (-x) as u32 } else { x as u32 };
+        loop {
+            b[p] = (x % 10) as u8 + b'0';
+            p -= 1;
+            x /= 10;
+            if x == 0 {
+                break;
+            }
+        }
+        if neg {
+            b[p] = b'-';
+            p -= 1;
+        }
+        if i > 0 {
+            b[p] = b'\n';
+            p -= 1;
+        }
     }
-    print!("{}", buffer);
+    p += 1;
+
+    stdout().lock().write(&b[p..]).unwrap();
 }
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);
         });
     }
 }
diff --git a/kattis-kth-alginda-quicksort/test b/kattis-kth-alginda-quicksort/test
index 4314cbf..01b0ae7 100755
--- a/kattis-kth-alginda-quicksort/test
+++ b/kattis-kth-alginda-quicksort/test
@@ -4,6 +4,7 @@ echo hej > out_1
 echo hej > out_2
 
 while diff out_1 out_2; do
+    echo -n .
     python gen.py -2147483648 2147483647 > max
     "$1" < max > out_1
     "$2" < max > out_2