From f621e603f33fab527d686bc110be8f212f3d8109 Mon Sep 17 00:00:00 2001
From: Antoine Martin <antoine@alarsyo.net>
Date: Tue, 7 Dec 2021 15:22:30 +0100
Subject: [PATCH] 2021: day07: QuickSelect for median computation

---
 Cargo.lock           | 64 ++++++++++++++++++++++++++++++++++++++++++++
 aoc2021/Cargo.toml   |  1 +
 aoc2021/src/day07.rs | 63 ++++++++++++++++++++++++++++++++++++++++---
 3 files changed, 125 insertions(+), 3 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index acee5e3..82780e0 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -88,6 +88,7 @@ version = "0.1.0"
 dependencies = [
  "anyhow",
  "aoc",
+ "rand",
 ]
 
 [[package]]
@@ -348,6 +349,17 @@ dependencies = [
  "typenum",
 ]
 
+[[package]]
+name = "getrandom"
+version = "0.2.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7fcd999463524c52659517fe2cea98493cfe485d10565e7b0fb07dbba7ad2753"
+dependencies = [
+ "cfg-if",
+ "libc",
+ "wasi",
+]
+
 [[package]]
 name = "half"
 version = "1.8.2"
@@ -528,6 +540,12 @@ dependencies = [
  "plotters-backend",
 ]
 
+[[package]]
+name = "ppv-lite86"
+version = "0.2.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ed0cfbc8191465bed66e1718596ee0b0b35d5ee1f41c5df2189d0fe8bde535ba"
+
 [[package]]
 name = "proc-macro2"
 version = "1.0.32"
@@ -552,6 +570,46 @@ version = "0.5.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "941ba9d78d8e2f7ce474c015eea4d9c6d25b6a3327f9832ee29a4de27f91bbb8"
 
+[[package]]
+name = "rand"
+version = "0.8.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2e7573632e6454cf6b99d7aac4ccca54be06da05aca2ef7423d22d27d4d4bcd8"
+dependencies = [
+ "libc",
+ "rand_chacha",
+ "rand_core",
+ "rand_hc",
+]
+
+[[package]]
+name = "rand_chacha"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88"
+dependencies = [
+ "ppv-lite86",
+ "rand_core",
+]
+
+[[package]]
+name = "rand_core"
+version = "0.6.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d34f1408f55294453790c48b2f1ebbb1c5b4b7563eb1f418bcfcfdbb06ebb4e7"
+dependencies = [
+ "getrandom",
+]
+
+[[package]]
+name = "rand_hc"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d51e9f596de227fda2ea6c84607f5558e196eeaf43c986b724ba4fb8fdf497e7"
+dependencies = [
+ "rand_core",
+]
+
 [[package]]
 name = "rayon"
 version = "1.5.1"
@@ -749,6 +807,12 @@ dependencies = [
  "winapi-util",
 ]
 
+[[package]]
+name = "wasi"
+version = "0.10.2+wasi-snapshot-preview1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fd6fbd9a79829dd1ad0cc20627bf1ed606756a7f77edff7b66b7064f9cb327c6"
+
 [[package]]
 name = "wasm-bindgen"
 version = "0.2.78"
diff --git a/aoc2021/Cargo.toml b/aoc2021/Cargo.toml
index b336e63..4d8f362 100644
--- a/aoc2021/Cargo.toml
+++ b/aoc2021/Cargo.toml
@@ -9,6 +9,7 @@ edition = "2021"
 [dependencies]
 aoc = { path = "../aoc" }
 anyhow = "1.0"
+rand = "0.8"
 
 [lib]
 path = "src/lib.rs"
diff --git a/aoc2021/src/day07.rs b/aoc2021/src/day07.rs
index 764df93..34151b8 100644
--- a/aoc2021/src/day07.rs
+++ b/aoc2021/src/day07.rs
@@ -1,6 +1,7 @@
 use std::fmt::Write;
 
 use anyhow::{Context, Result};
+use rand::Rng;
 
 const INPUT: &str = include_str!("../input/day07.txt");
 
@@ -20,9 +21,8 @@ fn part1(input: &str) -> Result<u64> {
         .map(|n| n.parse::<u64>().context("couldn't parse position"))
         .collect::<Result<Vec<_>>>()?;
 
-    // TODO: try linear selection algorithm
-    horizontal_positions.sort_unstable();
-    let median = horizontal_positions[horizontal_positions.len() / 2];
+    let median_rank = horizontal_positions.len() / 2;
+    let median = selection(&mut horizontal_positions, median_rank);
 
     Ok(horizontal_positions
         .iter()
@@ -31,6 +31,52 @@ fn part1(input: &str) -> Result<u64> {
         .sum())
 }
 
+fn selection<T>(data: &mut [T], i: usize) -> T
+where
+    T: Copy + Ord,
+{
+    if data.len() == 1 {
+        return data[0];
+    }
+
+    let mid = random_partition(data);
+
+    if i < mid {
+        selection(&mut data[..mid], i)
+    } else {
+        selection(&mut data[mid..], i - mid)
+    }
+}
+
+fn random_partition<T>(data: &mut [T]) -> usize
+where
+    T: Copy + Ord,
+{
+    let pivot_index = rand::thread_rng().gen_range(0..data.len());
+    let pivot = data[pivot_index];
+
+    let mut i = 0;
+    let mut j = data.len() - 1;
+
+    loop {
+        while data[i] < pivot {
+            i += 1;
+        }
+
+        while data[j] > pivot {
+            j -= 1;
+        }
+
+        if i >= j {
+            return usize::max(i, 1);
+        }
+
+        data.swap(i, j);
+        i += 1;
+        j -= 1;
+    }
+}
+
 fn part2(input: &str) -> Result<u64> {
     let horizontal_positions = input
         .trim()
@@ -90,4 +136,15 @@ mod tests {
     fn part2_real() {
         assert_eq!(part2(INPUT).unwrap(), 96592275);
     }
+
+    #[test]
+    fn test_selection() {
+        for _ in 0..4200 {
+            for i in 0..=9 {
+                let mut data = vec![9, 2, 7, 3, 5, 4, 6, 1, 8, 0];
+                let res = selection(&mut data, i);
+                assert_eq!(res, i);
+            }
+        }
+    }
 }