From 8017827edcff6df4576d4d77861e9735f261338a Mon Sep 17 00:00:00 2001 From: Antoine Martin Date: Mon, 14 Dec 2020 14:02:04 +0100 Subject: [PATCH] 2020: day14: part 2 --- ...day14_provided.txt => day14_provided1.txt} | 0 aoc2020/input/day14_provided2.txt | 4 + aoc2020/src/day14.rs | 149 +++++++++++++++++- 3 files changed, 151 insertions(+), 2 deletions(-) rename aoc2020/input/{day14_provided.txt => day14_provided1.txt} (100%) create mode 100644 aoc2020/input/day14_provided2.txt diff --git a/aoc2020/input/day14_provided.txt b/aoc2020/input/day14_provided1.txt similarity index 100% rename from aoc2020/input/day14_provided.txt rename to aoc2020/input/day14_provided1.txt diff --git a/aoc2020/input/day14_provided2.txt b/aoc2020/input/day14_provided2.txt new file mode 100644 index 0000000..bfcd365 --- /dev/null +++ b/aoc2020/input/day14_provided2.txt @@ -0,0 +1,4 @@ +mask = X1001X +mem[42] = 100 +mask = 00X0XX +mem[26] = 1 diff --git a/aoc2020/src/day14.rs b/aoc2020/src/day14.rs index 388fdda..fc5be3f 100644 --- a/aoc2020/src/day14.rs +++ b/aoc2020/src/day14.rs @@ -9,6 +9,7 @@ pub fn run() -> aoc::Result { let mut res = String::with_capacity(128); writeln!(res, "part 1: {}", part1(INPUT)?)?; + writeln!(res, "part 2: {}", part2(INPUT)?)?; Ok(res) } @@ -21,6 +22,14 @@ fn part1(input: &str) -> aoc::Result { Ok(program.memory_sum()) } +fn part2(input: &str) -> aoc::Result { + let mut program: Program = input.parse()?; + + program.run_part2()?; + + Ok(program.memory_sum()) +} + #[derive(Debug, Clone, Copy)] enum Mask { Floating, @@ -28,6 +37,106 @@ enum Mask { Zero, } +#[derive(Debug)] +struct FloatingIterator<'a> { + masks: &'a [Mask], + current_value: usize, + done: bool, + + stack: Vec<(usize, FloatingState)>, +} + +impl<'a> FloatingIterator<'a> { + fn new(masks: &'a [Mask], mut n: usize) -> Self { + for (offset, mask) in masks.iter().enumerate() { + if let Mask::One = mask { + n |= 1 << offset; + } + } + + Self { + masks, + current_value: n, + done: false, + stack: Vec::new(), + } + } + + fn find_next_floating(masks: &[Mask], offset: usize) -> Option { + masks + .iter() + .enumerate() + .skip(offset) + .find(|(_, m)| matches!(m, Mask::Floating)) + .map(|(idx, _)| idx) + } +} + +#[derive(Debug)] +enum FloatingState { + Unapplied, + AppliedZero, + AppliedZeroAndOne, +} + +impl<'a> Iterator for FloatingIterator<'a> { + type Item = usize; + + fn next(&mut self) -> Option { + if self.done { + return None; + } + + if self.stack.is_empty() { + // initialize stack with first floating iterator + let next = FloatingIterator::find_next_floating(self.masks, 0); + + match next { + Some(offset) => self.stack.push((offset, FloatingState::Unapplied)), + None => { + // there are no floating bits in this mask, we can return the value directly + self.done = true; + return Some(self.current_value); + } + } + } + + loop { + let (offset, state) = self.stack.last_mut().unwrap(); + match state { + FloatingState::Unapplied => { + self.current_value &= !(1 << *offset); + *state = FloatingState::AppliedZero; + } + FloatingState::AppliedZero => { + self.current_value |= 1 << *offset; + *state = FloatingState::AppliedZeroAndOne; + } + FloatingState::AppliedZeroAndOne => { + self.stack.pop(); + + if self.stack.is_empty() { + // we've computed all possibilities, we can just stop now + self.done = true; + return None; + } + + continue; + } + } + + // we've applied our current mask transform, now we "recurse" and find the next one + match FloatingIterator::find_next_floating(self.masks, *offset + 1) { + Some(offset) => self.stack.push((offset, FloatingState::Unapplied)), + None => { + // we were the last Floating mask, we can return the produced value + return Some(self.current_value); + } + } + } + } +} + #[derive(Debug, Clone)] struct BitMask { masks: Vec, @@ -46,6 +155,10 @@ impl BitMask { n } + + fn apply(&self, n: usize) -> impl Iterator + '_ { + FloatingIterator::new(&self.masks, n) + } } impl std::str::FromStr for BitMask { @@ -146,6 +259,27 @@ impl Program { Ok(()) } + fn run_part2(&mut self) -> aoc::Result<()> { + for inst in &self.instructions { + match inst { + Instruction::ChangeMask(bitmask) => self.current_mask = Some(bitmask.clone()), + + Instruction::MemWrite { offset, value } => match &self.current_mask { + Some(bitmask) => { + for offset in bitmask.apply(*offset) { + self.memory.insert(offset, *value); + } + } + None => { + return Err(err!("tried to execute MemWrite but mask isn't initialized")) + } + }, + } + } + + Ok(()) + } + fn memory_sum(&self) -> u64 { self.memory.iter().map(|(_, value)| value).sum() } @@ -172,15 +306,26 @@ impl std::str::FromStr for Program { mod tests { use super::*; - const PROVIDED: &str = include_str!("../input/day14_provided.txt"); + const PROVIDED1: &str = include_str!("../input/day14_provided1.txt"); + const PROVIDED2: &str = include_str!("../input/day14_provided2.txt"); #[test] fn part1_provided() { - assert_eq!(part1(PROVIDED).unwrap(), 165); + assert_eq!(part1(PROVIDED1).unwrap(), 165); } #[test] fn part1_real() { assert_eq!(part1(INPUT).unwrap(), 4297467072083); } + + #[test] + fn part2_provided() { + assert_eq!(part2(PROVIDED2).unwrap(), 208); + } + + #[test] + fn part2_real() { + assert_eq!(part2(INPUT).unwrap(), 5030603328768); + } }