aoc 2023 day 09

main.rs

use std::{io::BufRead, ops::ControlFlow};

fn main() {
    let mut context = Context::default();

    let mut s = String::new();

    let mut input = std::io::stdin().lock();

    let mut nums = Vec::new();

    let mut sum_of_predicted_forwards = 0;
    let mut sum_of_predicted_backwards = 0;

    loop {
        s.clear();
        let read = input.read_line(&mut s).unwrap();
        if read == 0 {
            break;
        }

        let line = s.trim();
        nums.clear();
        nums.extend(line.split_whitespace().map(|x| x.parse::<i64>().unwrap()));

        sum_of_predicted_forwards += dbg!(predict(&nums, &mut context, false));
        sum_of_predicted_backwards += dbg!(predict(&nums, &mut context, true));

        assert_eq!(context.run_ends.len(), 0);
        assert_eq!(context.delta_scratch.len(), 0, "leftovers with {nums:?}");
        assert_eq!(context.next_delta.len(), 0);
    }

    assert_eq!(sum_of_predicted_forwards, 1995001648);
    assert_eq!(sum_of_predicted_backwards, 988);
}

#[derive(Default, Debug)]
struct Context {
    delta_scratch: Vec<i64>,
    run_ends: Vec<usize>,
    next_delta: Vec<i64>,
}

impl Context {
    fn prev_by_ends(&self) -> &[i64] {
        &self.delta_scratch[self.prev_by_ends_range()]
    }

    fn prev_by_ends_range(&self) -> std::ops::Range<usize> {
        let prev = self.run_ends.len() - 2;
        let prev = self.run_ends[prev];
        let last = self.run_ends.len() - 1;
        let last = self.run_ends[last];

        prev..last
    }

    fn last_run_was_equal(&self) -> bool {
        matches!(
            self.prev_by_ends().iter().try_fold(None, |first, next| {
                if let Some((first, _)) = first {
                    if next != first {
                        return ControlFlow::Break(Some((next, false)));
                    }
                }
                ControlFlow::Continue(Some((next, true)))
            }),
            ControlFlow::Continue(Some((_, true)))
        )
    }

    fn push_delta_scratch_end(&mut self) {
        self.run_ends.push(self.delta_scratch.len());
    }

    fn init<I: Iterator<Item = i64>>(&mut self, deltas: I) {
        self.push_delta_scratch_end();
        self.delta_scratch.extend(deltas);
        self.push_delta_scratch_end();
    }

    fn delta_within_last_run(&mut self) {
        self.next_delta.clear();

        let scratch_range = self.prev_by_ends_range();

        let scratch = &self.delta_scratch[scratch_range];

        self.next_delta
            .extend(scratch.windows(2).map(|x| x[1] - x[0]));

        self.delta_scratch.append(&mut self.next_delta);

        self.push_delta_scratch_end();
    }

    fn collapse_prediction(&mut self) -> Option<i64> {
        if self.run_ends.len() < 2 {
            assert_eq!(
                self.run_ends.len(),
                1,
                "{:?} supposed to be have one zero",
                self.run_ends
            );

            return None;
        }

        let scratch_range = self.prev_by_ends_range();

        let scratch = self.delta_scratch.drain(scratch_range);
        let last = scratch.last().unwrap();

        // range to here has now been processed to last
        self.run_ends.pop().unwrap();

        let last_level_above = if self.run_ends.len() < 2 {
            0
        } else {
            self.delta_scratch[self.prev_by_ends_range()]
                .last()
                .copied()
                .unwrap()
        };

        if !self.delta_scratch.is_empty() {
            // this would just be junk; this needs to be summed up with the outer nums
            self.delta_scratch.push(last + last_level_above);
            self.run_ends.pop().unwrap();
            self.run_ends.push(self.delta_scratch.len());
        }

        Some(last)
    }

    fn collapse_predict_all(&mut self) -> i64 {
        let mut last = None;

        while let Some(pred) = self.collapse_prediction() {
            last = Some(pred);
        }

        assert_eq!(self.run_ends.pop().unwrap(), 0);

        last.unwrap()
    }
}

fn predict(nums: &[i64], context: &mut Context, backwards: bool) -> i64 {
    if !backwards {
        context.init(nums.windows(2).map(|slice| slice[1] - slice[0]));
    } else {
        context.init(nums.windows(2).rev().map(|slice| slice[1] - slice[0]))
    }

    while !context.last_run_was_equal() {
        // need to go one level deeper
        context.delta_within_last_run();
    }

    let collapsed = context.collapse_predict_all();

    if !backwards {
        collapsed + nums.last().unwrap()
    } else {
        -collapsed + nums.first().unwrap()
    }
}

#[test]
fn predict_example_1() {
    let example = [0, 3, 6, 9, 12, 15];
    let mut ctx = Context::default();
    let pred = predict(&example, &mut ctx, false);
    assert_eq!(pred, 18);
    let pred = predict(&example, &mut ctx, true);
    assert_eq!(pred, -3);
}

#[test]
fn predict_example_2() {
    let example = [1, 3, 6, 10, 15, 21];
    let mut ctx = Context::default();
    let pred = predict(&example, &mut ctx, false);
    assert_eq!(pred, 28);
    let pred = predict(&example, &mut ctx, true);
    assert_eq!(pred, 0);
}

#[test]
fn real_input_line_leftovers() {
    let example = [
        14, 16, 18, 20, 22, 24, 26, 28, 30, 32, 34, 36, 38, 40, 42, 44, 46, 48, 50, 52, 54,
    ];
    let mut ctx = Context::default();
    let _pred = predict(&example, &mut ctx, false);
    assert_eq!(ctx.run_ends.len(), 0);
    assert_eq!(ctx.delta_scratch.len(), 0, "leftovers: {ctx:?}");
    assert_eq!(ctx.next_delta.len(), 0);

    let _pred = predict(
        &[
            12, 27, 47, 68, 92, 149, 343, 943, 2562, 6504, 15418, 34500, 73686, 151710, 303812,
            596709, 1157888, 2232401, 4290657, 8227316, 15717112,
        ],
        &mut ctx,
        false,
    );
    assert_eq!(ctx.run_ends.len(), 0);
    assert_eq!(ctx.delta_scratch.len(), 0, "leftovers: {ctx:?}");
    assert_eq!(ctx.next_delta.len(), 0);
}