Day 14 Code

use std::env;
use std::io::{self};
use std::collections::HashMap;

extern crate regex;
use regex::Regex;

extern crate itertools;
use itertools::Itertools;

const PART1_STEPS: usize = 10;
const PART2_STEPS: usize = 40;
const DEBUG: bool = false;

// Part 1 (naive solution)
fn part1(poly: &String, rules: &HashMap<(char,char),char>) -> usize {
    let mut poly = poly.to_string();
    let mut diff = 0;
    for step in 0..PART1_STEPS {
        if DEBUG { println!("After step {}: {}", &step, &poly); }

        // Build new polymer
        let mut poly_new = String::new();
        for (l,r) in poly.chars().tuple_windows() {
            let rule = rules.get(&(l,r)).unwrap();
            poly_new.push_str(&format!("{}{}",l,rule));
        }
        poly_new.push_str(&format!("{}",poly.chars().last().unwrap()));
        poly = poly_new.clone();

        // Calculate metric
        let mut char_counts: HashMap<char,usize> = HashMap::new();
        for ch in poly_new.chars() {
            *char_counts.entry(ch).or_insert(0) += 1;
        }
        let most = char_counts.iter().map(|(_,v)| v).max().unwrap();
        let least = char_counts.iter().map(|(_,v)| v).min().unwrap();
        diff = most - least;
    }
    diff
}

// Part 2 (character counting)
fn part2(poly: &str, rules: &HashMap<(char,char),char>) -> usize {

    // Build initial counts of each letter and pair of letters
    let mut pair_counts: HashMap<(char,char),usize> = HashMap::new();
    let mut char_counts: HashMap<char,usize> = HashMap::new();
    for (l,r) in poly.chars().tuple_windows() {
        *pair_counts.entry((l,r)).or_insert(0) += 1;
    }
    for ch in poly.chars() {
        *char_counts.entry(ch).or_insert(0) += 1;
    }

    // Fold polymer
    for _ in 0..PART2_STEPS {
        let pairs_before = pair_counts.clone();
        pair_counts.clear();
        for (pair,count) in pairs_before.iter() {
            let insert = rules.get(&pair).unwrap();
            char_counts.entry(*insert).and_modify(|e| *e += count).or_insert(*count);
            pair_counts.entry((pair.0,*insert)).and_modify(|e| *e += count).or_insert(*count);
            pair_counts.entry((*insert,pair.1)).and_modify(|e| *e += count).or_insert(*count);
        }
    }
    let most = char_counts.iter().map(|(_,v)| v).max().unwrap();
    let least = char_counts.iter().map(|(_,v)| v).min().unwrap();
    most - least
}

fn solve(input: &str) -> io::Result<()> {
    // Input
    let input_str = std::fs::read_to_string(input).unwrap();
    let input_str = input_str.trim();
    let input: Vec<_> = input_str.split("\n\n").collect();

    // Initial polymer
    let poly = input[0].to_string();

    // Rules input
    let re = Regex::new(r"(\w)(\w) -> (\w)").unwrap();
    let rules: HashMap<(char,char),char> = input[1].split("\n")
        .map(|line| {
            let matches = re.captures(line).unwrap();
            let left   = matches[1].chars().next().unwrap();
            let right  = matches[2].chars().next().unwrap();
            let insert = matches[3].chars().next().unwrap();
            ((left,right),insert)
        })
        .collect();

    // Answers
    let part1 = part1(&poly, &rules);
    let part2 = part2(&poly, &rules);
    println!("Part 1: {}", part1); // 2194
    println!("Part 2: {}", part2); // 2360298895777

    Ok(())
}

fn main() {
    let args: Vec<String> = env::args().collect();
    let filename = &args[1];
    solve(&filename).unwrap();
}