Reimplementation of day 22 that doesn't need the whole deck in memory

1 files changed, 142 insertions, 72 deletions

diff --git a/src/bin/day_22.rs b/src/bin/day_22.rs
index cadfac4..fbe3f4c 100644
--- a/src/bin/day_22.rs
+++ b/src/bin/day_22.rs
@@ -1,13 +1,9 @@
-use cached::cached;
-use rpds::rbt_set;
-use rpds::vector::Vector;
-use rpds::RedBlackTreeMap;
-use rpds::RedBlackTreeMapSync;
+use num::bigint::BigInt;
+use num::traits::identities::Zero;
+use num::traits::sign::abs;
 use std::fmt;
 use std::io;
 use std::io::prelude::*;
-use std::iter;
-use std::iter::FromIterator;
 use std::process;
 use std::str::FromStr;
 use structopt::StructOpt;
@@ -19,9 +15,15 @@ use structopt::StructOpt;
 /// See https://adventofcode.com/2019/day/22 for details.
 struct Opt {
     /// The size of the deck
-    deck_size: usize,
+    deck_size: BigInt,
     /// At the end, query the position of card
-    card: usize,
+    card: BigInt,
+    /// Number of repetitions
+    repetitions: usize,
+
+    /// Prints the card in position n, rather than the position of card n
+    #[structopt(short = "p")]
+    position_mode: bool,
 }
 
 fn main() {
@@ -37,12 +39,27 @@ fn main() {
 
     //eprintln!("{:?}", instructions);
 
-    println!(
-        "{}",
-        Deck::new(opt.deck_size)
-            .shuffle(&instructions)
-            .find_position(opt.card)
-    );
+    if opt.position_mode {
+        println!(
+            "{}",
+            Card {
+                position: opt.card,
+                modulus: opt.deck_size
+            }
+            .shuffle_reverse(&instructions, opt.repetitions)
+            .position
+        );
+    } else {
+        println!(
+            "{}",
+            Card {
+                position: opt.card,
+                modulus: opt.deck_size
+            }
+            .shuffle(&instructions, opt.repetitions)
+            .position
+        );
+    }
 }
 
 fn exit_on_failed_assertion<A, E: std::error::Error>(data: Result<A, E>, message: &str) -> A {
@@ -56,91 +73,128 @@ fn exit_on_failed_assertion<A, E: std::error::Error>(data: Result<A, E>, message
 }
 
 #[derive(Clone)]
-struct Deck {
-    zero: usize,
-    cards: Vector<usize>,
+struct Card {
+    position: BigInt,
+    modulus: BigInt,
 }
 
-impl Deck {
-    fn new(deck_size: usize) -> Self {
-        Deck {
-            zero: 0,
-            cards: (0..deck_size).collect(),
-        }
-    }
-
-    fn shuffle(&self, instructions: &[Instruction]) -> Deck {
-        instructions.iter().fold(self.clone(), |deck, instruction| {
-            //eprintln!("{} :: {}", deck.zero, deck.cards);
-            deck.single_shuffle(instruction)
+impl Card {
+    fn shuffle(&self, instructions: &[Instruction], repetitions: usize) -> Card {
+        (0..repetitions).fold(self.clone(), |acc, _| {
+            instructions.iter().fold(acc, |card, instruction| {
+                //eprintln!("{} :: {}", deck.zero, deck.cards);
+                card.single_shuffle(instruction)
+            })
         })
     }
 
-    fn single_shuffle(&self, instruction: &Instruction) -> Deck {
+    fn single_shuffle(&self, instruction: &Instruction) -> Card {
         match instruction {
-            Instruction::DealIntoNewStack => Deck {
-                cards: (0..self.cards.len())
-                    .map(|i| self.at(self.cards.len() - i - 1))
-                    .collect(),
-                zero: 0,
+            Instruction::DealIntoNewStack => Card {
+                position: self.modulus.clone() - self.position.clone() - 1,
+                modulus: self.modulus.clone(),
             },
-            Instruction::Cut(n) => Deck {
-                zero: mod_plus(self.zero, *n, self.cards.len()),
-                ..self.clone()
+            Instruction::Cut(n) => Card {
+                position: mod_sub(self.position.clone(), n.clone(), self.modulus.clone()),
+                modulus: self.modulus.clone(),
             },
-            Instruction::ReverseCut(n) => Deck {
-                zero: mod_sub(self.zero, *n, self.cards.len()),
-                ..self.clone()
+            Instruction::ReverseCut(n) => Card {
+                position: mod_plus(self.position.clone(), n.clone(), self.modulus.clone()),
+                modulus: self.modulus.clone(),
             },
-            Instruction::DealWithIncrement(n) => Deck {
-                cards: (0..self.cards.len())
-                    .map(|i| self.at(*layout_after_shift(*n, self.cards.len()).get(&i).unwrap()))
-                    .collect(),
-                zero: 0,
+            Instruction::DealWithIncrement(n) => Card {
+                position: mod_times(self.position.clone(), n.clone(), self.modulus.clone()),
+                modulus: self.modulus.clone(),
             },
         }
     }
 
-    fn at(&self, index: usize) -> usize {
-        self.cards[mod_plus(index, self.zero, self.cards.len())]
+    fn shuffle_reverse(&self, instructions: &[Instruction], repetitions: usize) -> Card {
+        (0..repetitions).fold(self.clone(), |acc, _| {
+            instructions.iter().rev().fold(acc, |card, instruction| {
+                //eprintln!("{} :: {}", deck.zero, deck.cards);
+                card.single_shuffle_reverse(instruction)
+            })
+        })
     }
 
-    fn find_position(&self, value: usize) -> usize {
-        mod_sub(
-            self.cards.iter().position(|x| *x == value).unwrap(),
-            self.zero,
-            self.cards.len(),
-        )
+    fn single_shuffle_reverse(&self, instruction: &Instruction) -> Card {
+        match instruction {
+            Instruction::DealIntoNewStack => Card {
+                position: self.modulus.clone() - self.position.clone() - 1,
+                modulus: self.modulus.clone(),
+            },
+            Instruction::Cut(n) => Card {
+                position: mod_plus(self.position.clone(), n.clone(), self.modulus.clone()),
+                modulus: self.modulus.clone(),
+            },
+            Instruction::ReverseCut(n) => Card {
+                position: mod_sub(self.position.clone(), n.clone(), self.modulus.clone()),
+                modulus: self.modulus.clone(),
+            },
+            Instruction::DealWithIncrement(n) => Card {
+                position: mod_divide(self.position.clone(), n.clone(), self.modulus.clone()),
+                modulus: self.modulus.clone(),
+            },
+        }
     }
 }
 
-fn mod_plus(a: usize, b: usize, modulus: usize) -> usize {
+fn mod_plus(a: BigInt, b: BigInt, modulus: BigInt) -> BigInt {
     (a + b) % modulus
 }
 
-fn mod_sub(a: usize, b: usize, modulus: usize) -> usize {
-    (a + modulus - b) % modulus
+fn mod_sub(a: BigInt, b: BigInt, modulus: BigInt) -> BigInt {
+    (a + modulus.clone() - b) % modulus
 }
 
-fn mod_times(a: usize, b: usize, modulus: usize) -> usize {
+fn mod_times(a: BigInt, b: BigInt, modulus: BigInt) -> BigInt {
     (a * b) % modulus
 }
 
-cached! {
-    LAYOUT_AFTER_SHIFT;
-    fn layout_after_shift(n: usize, modulus: usize) -> RedBlackTreeMapSync<usize, usize> = {
-        (0..modulus).fold(RedBlackTreeMapSync::new_sync(), |acc, next| {
-            acc.insert(mod_times(next, n, modulus), next)
-        })
+fn mod_divide(a: BigInt, b: BigInt, modulus: BigInt) -> BigInt {
+    mod_times(a, mod_inverse(b, modulus.clone()), modulus)
+}
+
+// NB: This may give nonsense if modulus isn't coprime with a
+fn mod_inverse(a: BigInt, modulus: BigInt) -> BigInt {
+    (euclid_gcd_coefficients(a, modulus.clone()).0 + modulus.clone()) as BigInt % modulus
+}
+
+fn euclid_gcd_coefficients(a: BigInt, b: BigInt) -> (BigInt, BigInt) {
+    fn euclid_gcd_coefficients_inner(
+        r: BigInt,
+        old_r: BigInt,
+        s: BigInt,
+        old_s: BigInt,
+        t: BigInt,
+        old_t: BigInt,
+    ) -> (BigInt, BigInt) {
+        if r.is_zero() {
+            (old_s, old_t)
+        } else {
+            euclid_gcd_coefficients_inner(
+                old_r.clone() - (old_r.clone() / r.clone()) * r.clone(),
+                r.clone(),
+                old_s - (old_r.clone() / r.clone()) * s.clone(),
+                s,
+                old_t - (old_r.clone() / r) * t.clone(),
+                t,
+            )
+        }
     }
+
+    assert!(a < b);
+
+    euclid_gcd_coefficients_inner(b, a, 0.into(), 1.into(), 1.into(), 0.into())
 }
 
 #[derive(Debug)]
 enum Instruction {
     DealIntoNewStack,
-    Cut(usize),
-    ReverseCut(usize),
-    DealWithIncrement(usize),
+    Cut(BigInt),
+    ReverseCut(BigInt),
+    DealWithIncrement(BigInt),
 }
 
 impl FromStr for Instruction {
@@ -152,16 +206,16 @@ impl FromStr for Instruction {
             s.split(' ')
                 .nth(1)
                 .map(|val| {
-                    val.parse::<isize>()
+                    val.parse::<BigInt>()
                         .map_err(|_| ParseErr)
-                        .map(|parsed| Instruction::ReverseCut(parsed.abs() as usize))
+                        .map(|parsed| Instruction::ReverseCut(abs(parsed)))
                 })
                 .unwrap_or(Err(ParseErr))
         } else if s.starts_with("cut") {
             s.split(' ')
                 .nth(1)
                 .map(|val| {
-                    val.parse::<usize>()
+                    val.parse::<BigInt>()
                         .map_err(|_| ParseErr)
                         .map(|parsed| Instruction::Cut(parsed))
                 })
@@ -170,7 +224,7 @@ impl FromStr for Instruction {
             s.split(' ')
                 .nth(3)
                 .map(|val| {
-                    val.parse::<usize>()
+                    val.parse::<BigInt>()
                         .map_err(|_| ParseErr)
                         .map(|parsed| Instruction::DealWithIncrement(parsed))
                 })
@@ -191,3 +245,19 @@ impl fmt::Display for ParseErr {
 }
 
 impl std::error::Error for ParseErr {}
+
+#[test]
+fn mod_inverse_of_13() {
+    assert_eq!(mod_inverse(1.into(), 13.into()), 1.into());
+    assert_eq!(mod_inverse(2.into(), 13.into()), 7.into());
+    assert_eq!(mod_inverse(3.into(), 13.into()), 9.into());
+    assert_eq!(mod_inverse(4.into(), 13.into()), 10.into());
+    assert_eq!(mod_inverse(5.into(), 13.into()), 8.into());
+    assert_eq!(mod_inverse(6.into(), 13.into()), 11.into());
+    assert_eq!(mod_inverse(7.into(), 13.into()), 2.into());
+    assert_eq!(mod_inverse(8.into(), 13.into()), 5.into());
+    assert_eq!(mod_inverse(9.into(), 13.into()), 3.into());
+    assert_eq!(mod_inverse(10.into(), 13.into()), 4.into());
+    assert_eq!(mod_inverse(11.into(), 13.into()), 6.into());
+    assert_eq!(mod_inverse(12.into(), 13.into()), 12.into());
+}