Solution for Advent of Code 2017 Day 3

day03.rs

use std::time::Instant;

use itertools::Itertools;
use rustc_hash::FxHashMap;

const INPUT: &str = include_str!("inputs/day03.txt");

fn parse(input: &str) -> usize {
    input.trim().parse::<usize>().unwrap()
}

fn p1(input: &str) -> usize {
    let input = parse(input);

    let mut v = 1;
    let mut r = 0;
    while v * v < input {
        r += 1;
        v += 2;
    }

    let midpoints = [
        v * v - v / 2,
        v * v - v / 2 - v - 1,
        v * v - v / 2 - (v - 1) * 2,
        v * v - v / 2 - (v - 1) * 3,
    ];

    let nearest = midpoints.iter().map(|m| input.abs_diff(*m)).min().unwrap();
    r + nearest
}

fn p2(input: &str) -> usize {
    let input = parse(input);

    let mut seen = FxHashMap::default();
    seen.insert((0, 0), 1);
    seen.insert((1, 0), 1);

    let (mut cx, mut cy) = (1, 0);
    let mut steps = 1;
    let directions = [(0, 1), (-1, 0), (0, -1), (1, 0)];

    // Loop through all the directions until we've found the value.
    loop {
        for (i, (dx, dy)) in directions.iter().enumerate() {
            // We step longer in the directions left and right.
            if i == 1 || i == 3 {
                steps += 1;
            }
            for _ in 0..steps {
                cx += dx;
                cy += dy;

                // Calculate
                let cur = (-1..=1)
                    .cartesian_product(-1..=1)
                    .map(|(x, y)| seen.get(&(cx + x, cy + y)).unwrap_or(&0))
                    .sum::<usize>();

                if cur >= input {
                    return cur;
                }
                seen.insert((cx, cy), cur);
            }
        }
    }
}

fn main() {
    let now = Instant::now();
    let solution = p1(INPUT);
    println!("p1 {:?} {}", now.elapsed(), solution);

    let now = Instant::now();
    let solution = p2(INPUT);
    println!("p2 {:?} {}", now.elapsed(), solution);
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_p1() {
        let input = "1024";
        assert_eq!(p1(input), 31);
        let input = "1";
        assert_eq!(p1(input), 0);
        let input = "12";
        assert_eq!(p1(input), 3);
        let input = "23";
        assert_eq!(p1(input), 2);
    }

    #[test]
    fn test_p2() {
        let input = "800";
        assert_eq!(p2(input), 806);
    }
}