Faster approach only traversing tree as much as needed

Author: maneatingape

README.md

@@ -65,10 +65,10 @@ pie
     "Year 2016" : 663
     "Year 2020" : 286
     "Year 2015" : 87
-    "Years 2019, 2021 and 2022" : 47
+    "Years 2019, 2021 and 2022" : 45
 ```
 
-* [2022](#2022) (14 ms)
+* [2022](#2022) (12 ms)
 * [2021](#2021) (11 ms)
 * [2020](#2020) (286 ms)
 * [2019](#2019) (22 ms)
@@ -81,7 +81,7 @@ pie
 %%{init: {"themeVariables": { "pie1": "#7cb5ec", "pie2": "#90ed7d", "pie3": "#f7a35c", "pie4": "#c6cdd5"}}}%%
 pie
     title Year 2022 benchmark time by % of total
-    "Day 20" : 7449
+    "Day 20" : 5067
     "Day 11" : 2386
     "Day 23" : 2017
     "Others" : 2029
@@ -108,7 +108,7 @@ pie
 | 17 | [Pyroclastic Flow](https://adventofcode.com/2022/day/17) | [Source](src/year2022/day17.rs) | 68 |
 | 18 | [Boiling Boulders](https://adventofcode.com/2022/day/18) | [Source](src/year2022/day18.rs) | 129 |
 | 19 | [Not Enough Minerals](https://adventofcode.com/2022/day/19) | [Source](src/year2022/day19.rs) | 407 |
-| 20 | [Grove Positioning System](https://adventofcode.com/2022/day/20) | [Source](src/year2022/day20.rs) | 7449 |
+| 20 | [Grove Positioning System](https://adventofcode.com/2022/day/20) | [Source](src/year2022/day20.rs) | 5067 |
 | 21 | [Monkey Math](https://adventofcode.com/2022/day/21) | [Source](src/year2022/day21.rs) | 61 |
 | 22 | [Monkey Map](https://adventofcode.com/2022/day/22) | [Source](src/year2022/day22.rs) | 132 |
 | 23 | [Unstable Diffusion](https://adventofcode.com/2022/day/23) | [Source](src/year2022/day23.rs) | 2017 |

src/year2022/day20.rs

@@ -2,21 +2,21 @@ use crate::util::iter::*;
 use crate::util::parse::*;
 use std::ops::{Index, IndexMut};
 
-#[derive(Clone, Copy)]
+#[derive(Clone, Copy, Default)]
 struct Node {
     size: u16,
     left: u16,
     right: u16,
     up: u16,
 }
 
-struct Tree {
+struct OrderStatisticTree {
     root: u16,
-    size: i64,
+    size: u16,
     nodes: Vec<Node>,
 }
 
-impl Index<u16> for Tree {
+impl Index<u16> for OrderStatisticTree {
     type Output = Node;
 
     #[inline]
@@ -25,27 +25,27 @@ impl Index<u16> for Tree {
     }
 }
 
-impl IndexMut<u16> for Tree {
+impl IndexMut<u16> for OrderStatisticTree {
     #[inline]
     fn index_mut(&mut self, index: u16) -> &mut Self::Output {
         &mut self.nodes[index as usize]
     }
 }
 
-impl Tree {
-    fn from(input: &[i64]) -> Tree {
-        let mut size = input.len();
+impl OrderStatisticTree {
+    fn from(size: usize) -> Self {
         let mut start = 0;
         let mut end = size;
+        let mut level = size;
 
-        let mut remaining = None;
         let mut nodes = Vec::with_capacity(size * 2);
+        let mut remaining = None;
 
         let empty = u16::MAX;
-        let leaf = Node { size: 1, left: empty, right: empty, up: empty };
+        let leaf = Node { size: 1, ..Default::default() };
         nodes.resize(size, leaf);
 
-        while size > 0 {
+        while level > 0 {
             let mut push = |left: usize, right: usize| {
                 let index = nodes.len() as u16;
                 nodes[left].up = index;
@@ -59,7 +59,7 @@ impl Tree {
                 push(left, right);
             }
 
-            if size % 2 == 1 {
+            if level % 2 == 1 {
                 if let Some(right) = remaining {
                     remaining = None;
                     push(end - 1, right);
@@ -70,13 +70,13 @@ impl Tree {
 
             start = end;
             end = nodes.len();
-            size = end - start;
+            level = end - start;
         }
 
-        Tree { root: (end - 1) as u16, size: (input.len() - 1) as i64, nodes }
+        OrderStatisticTree { root: (end - 1) as u16, size: (size - 1) as u16, nodes }
     }
 
-    fn position(&self, start: usize) -> u16 {
+    fn position(&self, start: usize) -> usize {
         let mut cur = start as u16;
         let mut offset = 0;
 
@@ -89,12 +89,12 @@ impl Tree {
             cur = next;
         }
 
-        offset
+        offset as usize
     }
 
-    fn value_at(&self, start: u16) -> usize {
+    fn value_at(&self, start: usize) -> usize {
         let mut cur = self.root;
-        let mut offset = start;
+        let mut offset = start as u16;
 
         loop {
             let Node { size, left, right, .. } = self[cur];
@@ -114,71 +114,116 @@ impl Tree {
     }
 
     fn mix(&mut self, start: usize, value: i64) {
-        let start = start as u16;
-        let mut cur = start;
-        let mut offset = 0;
+        let mut delta_right = value.rem_euclid(self.size as i64) as u16;
+        let mut delta_left = self.size - delta_right;
 
-        while cur != self.root {
-            let next = self[cur].up;
-            let Node { left, right, .. } = self[next];
-
-            if right == cur {
-                offset += self[left].size;
-            };
-            self[next].size -= 1;
-            cur = next;
+        if delta_left == 0 || delta_right == 0 {
+            return;
         }
 
+        // Temporarily set our size to zero
+        let start = start as u16;
+        self[start].size = 0;
+
         let parent = self[start].up;
+        let mut dest = parent;
 
-        if parent == self.root {
-            if self[self.root].left == start {
-                self.root = self[self.root].right;
-            } else {
-                self.root = self[self.root].left;
+        // Go up
+        let (left, right) = loop {
+            let Node { left, right, up, .. } = self[dest];
+            let left_size = self[left].size;
+            let right_size = self[right].size;
+
+            if delta_left <= left_size || delta_right <= right_size {
+                break (left, right);
             }
-        } else {
-            let grand_parent = self[parent].up;
-            let next_parent =
-                if self[parent].left == start { self[parent].right } else { self[parent].left };
-            self[next_parent].up = grand_parent;
-            if self[grand_parent].left == parent {
-                self[grand_parent].left = next_parent;
+
+            if self[up].left == dest {
+                delta_right -= right_size;
+                delta_left += right_size;
             } else {
-                self[grand_parent].right = next_parent;
+                delta_right += left_size;
+                delta_left -= left_size;
             }
-        }
 
-        cur = self.root;
-        offset = (offset as i64 + value).rem_euclid(self.size) as u16;
+            self[dest].size -= 1;
+            dest = up;
+        };
 
-        loop {
-            let Node { size, left, right, up } = self[cur];
+        let (left, right) = if delta_left <= self[left].size {
+            // Down left
+            dest = left;
 
-            if size == 1 {
-                self[parent] = Node { size: 2, left: start, right: cur, up };
+            loop {
+                let Node { size, left, right, .. } = self[dest];
 
-                self[cur].up = parent;
-                self[start].up = parent;
+                if size == 1 {
+                    break (start, dest);
+                }
 
-                if self[up].left == cur {
-                    self[up].left = parent;
+                self[dest].size += 1;
+                let right_size = self[right].size;
+
+                if delta_left <= right_size {
+                    dest = right;
                 } else {
-                    self[up].right = parent;
+                    delta_left -= right_size;
+                    dest = left;
                 }
+            }
+        } else {
+            // Down right
+            dest = right;
+
+            loop {
+                let Node { size, left, right, .. } = self[dest];
+
+                if size == 1 {
+                    break (dest, start);
+                }
+
+                self[dest].size += 1;
+                let left_size = self[left].size;
 
-                break;
+                if delta_right <= left_size {
+                    dest = left;
+                } else {
+                    delta_right -= left_size;
+                    dest = right;
+                }
             }
+        };
 
-            self[cur].size += 1;
-            let size = self[left].size;
-            if offset < size {
-                cur = left;
+        // Remove node
+        let sibling =
+            if self[parent].left == start { self[parent].right } else { self[parent].left };
+
+        if parent == self.root {
+            self[sibling].size += 1;
+            self.root = sibling;
+        } else {
+            let grand_parent = self[parent].up;
+            self[sibling].up = grand_parent;
+
+            if self[grand_parent].left == parent {
+                self[grand_parent].left = sibling;
             } else {
-                cur = right;
-                offset -= size;
+                self[grand_parent].right = sibling;
             }
         }
+
+        // Insert node
+        let up = self[dest].up;
+
+        if self[up].left == dest {
+            self[up].left = parent;
+        } else {
+            self[up].right = parent;
+        }
+
+        self[parent] = Node { size: 2, left, right, up };
+        self[left] = Node { size: 1, up: parent, ..Default::default() };
+        self[right] = Node { size: 1, up: parent, ..Default::default() };
     }
 }
 
@@ -195,7 +240,7 @@ pub fn part2(input: &[i64]) -> i64 {
 }
 
 fn decrypt(input: &[i64], key: i64, rounds: usize) -> i64 {
-    let mut tree = Tree::from(input);
+    let mut tree = OrderStatisticTree::from(input.len());
 
     for _ in 0..rounds {
         for (i, n) in input.iter().enumerate() {
@@ -207,7 +252,7 @@ fn decrypt(input: &[i64], key: i64, rounds: usize) -> i64 {
     let zeroth = tree.position(start);
     [1000, 2000, 3000]
         .iter()
-        .map(|offset| tree.value_at((zeroth + offset) % (input.len() as u16)))
+        .map(|offset| tree.value_at((zeroth + offset) % input.len()))
         .map(|index| input[index] * key)
         .sum()
 }