Day 14: Extended Polymerization

template = ls[0]
rules = ls[2:] > fe(X.split(" -> ") | (λ x: (tuple(x[0]), x[1]))) | set
chars = set(template) | set(collapse(rules))
pairs = list(product(chars, repeat=2))
state = np.zeros((len(pairs), 1), dtype=object)
for pair in pairwise(template):
    state[pairs.index(pair)] += 1
M = np.zeros((len(pairs), len(pairs)), dtype=object)
for pair, insert in rules:
    M[pairs.index((pair[0], insert)), pairs.index(pair)] = 1
    M[pairs.index((insert, pair[1])), pairs.index(pair)] = 1
# use (matrix) exponentiation by squaring to get O(log(n))
state = np.linalg.matrix_power(M, 40) @ state
# factor of 2 to avoid double-counting
# ceil to correct for template[0] and template[-1] off-by-one's
counts = [ceil(sum(count * pair.count(c)
                   for pair, count in zip(pairs, state[:, 0]))
               / 2)
          for c in chars]
max(counts) - min(counts)

#!/usr/bin/env python3
from itertools import pairwise
from itertools import product
from math import ceil
from more_itertools import collapse
from pipetools import X
from pipetools import foreach
import sys
from pyp import pypprint
fe = foreach
import numpy as np
lines = [x.rstrip('\n') for x in sys.stdin]
ls = lines
template = ls[0]
rules = ls[2:] > fe(X.split(' -> ') | (lambda x: (tuple(x[0]), x[1]))) | set
chars = set(template) | set(collapse(rules))
pairs = list(product(chars, repeat=2))
state = np.zeros((len(pairs), 1), dtype=object)
for pair in pairwise(template):
    state[pairs.index(pair)] += 1
M = np.zeros((len(pairs), len(pairs)), dtype=object)
for (pair, insert) in rules:
    M[pairs.index((pair[0], insert)), pairs.index(pair)] = 1
    M[pairs.index((insert, pair[1])), pairs.index(pair)] = 1
state = np.linalg.matrix_power(M, 40) @ state
counts = [ceil(sum((count * pair.count(c) for (pair, count) in zip(pairs, state[:, 0]))) / 2) for c in chars]
output = max(counts) - min(counts)
if output is not None:
    pypprint(output)

Then for the output

=ARRAYFORMULA(
  IF(C3=FALSE,,
   QUERY(
    QUERY(
     TRANSPOSE(
      SPLIT(
       REGEXREPLACE(E11,"(.)","$1|"),
       "|")),
     "select Col1, Count(Col1)
      where Col1 is not null
      group by Col1
      order by Count(Col1) desc
      label Count(Col1) ''"),
    "select Max(Col2)-Min(Col2)
     label Max(Col2)-Min(Col2) 'Part 1'")))

From there, change the reference to the one above and drag that shit down. I hate formulas like this, but I think its the only way.

I represent the rules as a dict of 2-tuples of characters mapping to their output character. that dovetails nicely with my favorite trick for walking consecutive pairs of an array, zip(polymer, polymer[1:])

biggest stumbling block I had was figuring out how to avoid duplicates in the middle of the array, but still tack on the final character at the end.

from collections import Counter
import re

RULE_RE = re.compile('(\w+) -> (\w+)')

def main():
    with open('014.txt') as file:
        lines = file.read().splitlines()

    template = lines[0]
    rules = {}
    for line in lines[2:]:
        match = RULE_RE.match(line)
        input, output = match.groups()
        a, b = input
        rules[(a, b)] = output

    polymer = template
    for step in range(10):
        polymer = expand(polymer, rules)

    counter = Counter(polymer)
    most = max(value for value in counter.values())
    least = min(value for value in counter.values())
    print(most - least)


def expand(polymer, rules):
    new = []

    for input_a, input_b in zip(polymer, polymer[1:]):
        new.append(input_a)

        for (a, b), output in rules.items():
            if input_a == a and input_b == b:
                new.append(output)
                break

    new.append(input_b)

    return ''.join(new)


if __name__ == '__main__':
    main()

I currently have an off by thirty error so I'm going to look at this tomorrow.

seriously. off by 30. take that, off-by-one errors. suck it.

I figured out an improved algorithm for part 2, tracking populations of pairs rather than a giant string. it seems to work.

I run it on 10 steps of the example input: Counter({'B': 1749, 'N': 865, 'C': 298, 'H': 161})

their example from part 1:

After step 10, B occurs 1749 times, C occurs 298 times, H occurs 191 times, and N occurs 865 times;

it's exactly right, except for H. H is off by 30.

edit: goddamnit. I am preserving this rant for posterity. it's a goddamn typo on their website.

After step 10, B occurs 1749 times, C occurs 298 times, H occurs 191 times, and N occurs 865 times; taking the quantity of the most common element (B, 1749) and subtracting the quantity of the least common element (H, 161) produces 1749 - 161 = 1588.

grumble grumble I want my money back. literally unplayable.

thinking about the previous "oh shit part 2 is going to take a completely infeasible amount of time and memory" problem (day 6) helped me here. with that, the trick was to take the thing we had eleventy gajillion individual objects of (fish) and turn them into a count instead. same thing works here.

I realized that rather than output characters in a string, this can be modeled as a pair of characters divides and creates two new pairs. and then I can simply track population counts.

I convert my rules dict to now keep track of the two pairs that should be generated when any existing pair divides. early on I had a stupid "mental copy-paste" bug here when I had [(a, output), (b, output)] instead of [(a, output), (output, b)].

the next stumbling block was "ok, I have population of pairs, great, how do I convert that into a population of individuals?".

the key trick I realized is that double-counting them as part of each pair is fine, because there's only two characters I don't double-count, and that's the very beginning and end of the string. and those two characters don't ever change from the original template. so I can double-count them manually and then when I divide the final result by 2 I get the exact answer.

from collections import Counter
import re

RULE_RE = re.compile('(\w+) -> (\w+)')

def main():
    with open('014.txt') as file:
        lines = file.read().splitlines()

    template = lines[0]
    rules = {}
    for line in lines[2:]:
        match = RULE_RE.match(line)
        input, output = match.groups()
        a, b = input
        rules[(a, b)] = [(a, output), (output, b)]

    population = Counter()
    for a, b in zip(template, template[1:]):
        population[(a, b)] += 1

    for i in range(40):
        population = expand(population, rules)

    individual_population = Counter()
    for (a, b), count in population.items():
        individual_population[a] += count
        individual_population[b] += count

    individual_population[template[0]] += 1
    individual_population[template[-1]] += 1

    most = max(value for value in individual_population.values())
    least = min(value for value in individual_population.values())
    diff = int((most - least) / 2)
    print(diff)


def expand(population, rules):
    new = Counter()

    for pair, count in population.items():
        for child in rules[pair]:
            new[child] += count

    return new


if __name__ == '__main__':
    main()

#!/usr/bin/awk -f
NR == 1 {
	str = $0

	split($0, chars, "")
	for (i in chars)
		count[chars[i]] += 1
}

/->/ {
	split($0, pair, " -> ")
	map[pair[1]] = pair[2]
}

END {
	for (step = 1; step <= 10; ++step) {
		for (i in map) {
			split(i, rule, "")
			while (index(str, i))
				count[map[i]] += gsub(i, rule[1] tolower(map[i]) rule[2], str)
		}

		str = toupper(str)
		print str
	}

	for (i in count) {
		if (count[i] > max || max == "")
			max = count[i]

		if (count[i] < min || min == "")
			min = count[i]
	}

	print max - min
}

#!/usr/bin/awk -f
NR == 1 {
	str = $0
	split($0, chars, "")
	for (i in chars)
		count[chars[i]] += 1
}

/->/ {
	split($0, rule, " -> ")
	map[rule[1]] = rule[2]
}

END {
	for (i in map) {
		split(i, rule, "")
		while (index(str, i))
			pair[i] += gsub(i, rule[1] "." rule[2], str)
	}

	for (step = 1; step <= 40; ++step) {
		for (i in pair)
			state[i] = pair[i]

		for (i in map) {
			split(i, rule, "")

			count[map[i]] += pair[i]
			state[i] -= pair[i]
			state[rule[1] map[i]] += pair[i]
			state[map[i] rule[2]] += pair[i]
		}

		for (i in state)
			pair[i] = state[i]
	}

	for (i in count) {
		if (count[i] > max || max == "")
			max = count[i]

		if (count[i] < min || min == "")
			min = count[i]
	}

	print max - min
}

def compute_p1(input)
  pairs, counts, rules = parse(input)

  10.times { apply_rules(pairs, counts, rules) }

  scores = counts.values.sort
  scores.last - scores.first
end

def apply_rules(pairs, counts, rules)
  pair_deltas = Hash.new(0)
  count_deltas = Hash.new(0)

  pairs.keys.filter{ pairs[_1] > 0 }.each do |key|
    count = pairs[key]
    pair_deltas[key] -= count
    rules[key][0..1].each { pair_deltas[_1] += count }
    count_deltas[rules[key].last] += count
  end

  pair_deltas.each { |k,v| pairs[k] += pair_deltas[k] }
  count_deltas.each { |k,v| counts[k] += count_deltas[k] }
end

def parse(input)
  template, rules = input.split("\n\n")

  rules = rules
            .lines
            .map(&:chomp)
            .map { _1.split(' -> ') }
            .map { |k,v| [k, [k[0] + v, v + k[1], v]] }
            .to_h

  pairs = template
            .chars
            .each_cons(2)
            .map(&:join)
            .reduce(Hash.new(0)) { |h, pair| h[pair] += 1; h }

  counts = template
             .chars
             .reduce(Hash.new(0)) { |h,c| h[c] += 1; h }

  [pairs, counts, rules]
end

With the optimization already done, we just run the loop a few more times.

def compute_p2(input)
  pairs, counts, rules = parse(input)

  40.times { apply_rules(pairs, counts, rules) }

  scores = counts.values.sort
  scores.last - scores.first
end

I took a bit of a gamble and assumed order wasn't going to matter for this, so I just looked at how one pair turns into 2 pairs and kept a count of how many of each pair I had. I added the elements after the fact because that seemed like the easiest way to figure out the individual elements. It basically rides on the back of the pairs calculation which is doing the real work

import re
from collections import defaultdict

#data = test_data_insertion
data = real_data_insertion
data = data.split('\n')

#starting_pattern = test_data_starting
starting_pattern = real_data_starting

pairs = defaultdict(int)
elements = defaultdict(int)
for element in list(starting_pattern):
    elements[element] = elements[element] + 1
for index,char in enumerate(list(starting_pattern)[:-1]):
    pair = char + list(starting_pattern)[index+1]
    pairs[pair] = pairs[pair] + 1

conversions = {}
convert = re.compile(r'((\w)(\w)) -> (\w)')
for row in data:
    results = convert.search(row)
    conversions[results[1]] = (results[2] + results[4], results[4]+results[3])
#print(conversions)

for step in range(10):
    new_pairs = defaultdict(int)
    new_elements = defaultdict(int)
    for key in pairs.keys():
        new_key_1 = conversions[key][0]
        new_key_2 = conversions[key][1]
        new_pairs[new_key_1] = new_pairs[new_key_1] + pairs[key]
        new_pairs[new_key_2] = new_pairs[new_key_2] + pairs[key]
        elements[new_key_1[1]] = elements[new_key_1[1]] + pairs[key]
    pairs = new_pairs
print(elements)
print(pairs)
print(max(elements.values())-min(elements.values()))

It's the same except for changing the run duration

• This is another efficiency problem, your number of elements almost, but not quite doubles every pass

• If you approach the problem certain ways, order doesn't end up being super important

I'm using 3 hashes stored together in a list.

• element -> count
• pair -> count
• pair -> element

(define (parse-polymer input)
  (list (count-elements (first input))
        (count-pairs (first input))
        (for/hash ([l (drop input 2)])
          (apply values (string-split l " -> ")))))

(define (count-elements p)
  (for/fold ([es (hash)])
            ([e (in-string p)])
    (hash-update es (string e) add1 0)))

(define (count-pairs template)
  (define es (string->list template))
  (for/fold ([pairs (hash)])
            ([e1 es] [e2 (rest es)])
    (hash-update pairs (string e1 e2) add1 0)))

To update the polymer template, I iterate through the element pairs, generating a new count of pairs along the way. For each pair:

• add its count for the newly inserted element.
• add its count for the 2 new element pairs made when adding the inserted element to the pair

(define (part-01 input)
  (insert-pairs input 10))

(define (insert-pairs input steps)
  (for/fold ([polymer (parse-polymer input)]
             #:result (polymer-score (first polymer)))
            ([_ steps])
    (apply update-polymer polymer)))

(define (polymer-score elements)
  (- (apply max (hash-values elements))
     (apply min (hash-values elements))))

(define (update-polymer es ps rs)
  (define ((inc a) b) (+ a b))
  (for*/fold ([es+ es]
              [ps+ (hash)]
              #:result (list es+ ps+ rs))
             ([(p pc) (in-hash ps)] #:unless (zero? pc))
    (define element (hash-ref rs p))
    (define e1 (string-append (substring p 0 1) element))
    (define e2 (string-append element (substring p 1)))
    (values (hash-update es+ element (inc pc) 0)
            (hash-update (hash-update ps+ e1 (inc pc) 0) e2 (inc pc) 0))))

lol

(define (part-02 input)
  (insert-pairs input 40))

from collections import Counter

polymer, _, *rules = open("input.txt").read().splitlines()
rules = dict(rule.split(" -> ") for rule in rules)
pairs = Counter(polymer[i] + polymer[i + 1] for i in range(len(polymer) - 1))
chars = Counter(polymer)
steps = (10, 40)

for i in range(1, steps[-1] + 1):
    for (a, b), count in pairs.copy().items():
        if a + b in rules.keys():
            insert = rules[a + b]
            pairs[a + b] -= count
            pairs[a + insert] += count
            pairs[insert + b] += count
            chars[insert] += count
    if i in steps:
        print("Answer:", max(chars.values()) - min(chars.values()))

This is the Lanternfish problem again, in different clothes.

Instead of counting single fish, we're counting pairs of characters. Each insertion AxB reduces the count of AB polymers to zero, and correspondingly increases the count of Ax and xB polymers.
The added twist is that you have to cache the updates, because the addition only takes effect after a round of substitutions is complete.

This was a Day 8 all over again for me. Simple problem, simple solution, but I botched it at some point and spent the better part of 1.5h chasing the bug. Unfortunately the test-cases were of no help, because my code passed them with flying colors -- it would just not produce an answer that the site would accept.

In desperation, I finally caved on my "only Rust" rule and hacked together a small Ruby script that did it by brute-force and then compared its output, step by painful step. Turns out that I was overwriting my character counts at one place instead of adding them; this only manifests if the same character pair appears twice, which is why the test input worked, but the real one didn't.

use std::collections::BTreeMap;

type Pair = (char, char);
type PairInsertionRule = (Pair, char);

#[derive(Debug, Clone)]
pub struct Polymerization {
    rules: Vec<PairInsertionRule>,
    template: String,
    counts: BTreeMap<Pair, isize>,
}

mod parse {
    use super::{BTreeMap, Pair, PairInsertionRule, Polymerization};

    use nom::{
        bytes::complete::tag,
        character::complete::{alpha1, digit1, line_ending},
        combinator::{eof, opt},
        multi::many1,
        IResult,
    };

    pub fn template(input: &str) -> IResult<&str, (String, BTreeMap<Pair, isize>)> {
        let (input, t) = alpha1(input)?;
        let (input, _) = line_ending(input)?;

        let mut result = BTreeMap::new();

        let mut chars: Vec<char> = t.chars().collect();
        chars.push('$');

        for w in chars.windows(2) {
            result
                .entry((w[0], w[1]))
                .and_modify(|cnt| *cnt += 1)
                .or_insert(1);
        }

        Ok((input, (t.to_string(), result)))
    }

    pub fn rule(input: &str) -> IResult<&str, PairInsertionRule> {
        let (input, start) = alpha1(input)?;
        let (input, _) = tag(" -> ")(input)?;
        let (input, insert) = alpha1(input)?;
        let (input, _) = line_ending(input)?;

        let mut chrs = start.chars();
        let a = chrs.next().unwrap();
        let b = chrs.next().unwrap();
        let c = insert.chars().next().unwrap();

        Ok((input, ((a, b), c)))
    }

    pub fn parse(input: &str) -> IResult<&str, Polymerization> {
        let (input, (t, template)) = template(input)?;
        let (input, _) = line_ending(input)?;
        let (input, rules) = many1(rule)(input)?;
        let (input, _) = eof(input)?;

        Ok((
            input,
            Polymerization {
                rules,
                template: t,
                counts: template,
            },
        ))
    }
}

pub fn pstep(p: &Polymerization) -> Polymerization {
    let mut next_counts = p.counts.clone();

    for &((a, b), c) in p.rules.iter() {
        let cur_count = *p.counts.get(&(a, b)).unwrap_or(&0);

        if cur_count == 0 {
            continue;
        }

        next_counts
            .entry((a, b))
            .and_modify(|cnt| *cnt -= cur_count)
            .or_insert(0);

        next_counts
            .entry((a, c))
            .and_modify(|cnt| *cnt += cur_count)
            .or_insert(cur_count);

        next_counts
            .entry((c, b))
            .and_modify(|cnt| *cnt += cur_count)
            .or_insert(cur_count);
    }

    Polymerization {
        rules: p.rules.clone(),
        template: p.template.clone(),
        counts: next_counts,
    }
}

fn count_occurrences(p: &Polymerization) -> BTreeMap<char, isize> {
    let mut cnts = BTreeMap::new();

    for ((a, _b), cnt) in p.counts.iter() {
        cnts.entry(*a).and_modify(|c| *c += *cnt).or_insert(*cnt);
    }

    cnts
}

fn part1(parsed: &Polymerization) -> isize {
    let result = (1..=10).fold(parsed.clone(), |p, _| {
        let s = pstep(&p);
        s
    });

    let cnts = count_occurrences(&result);

    let min = cnts.values().min().unwrap();
    let max = cnts.values().max().unwrap();

    max - min
}

fn part2(parsed: &Polymerization) -> isize {
    let result = (1..=40).fold(parsed.clone(), |p, _| {
        let s = pstep(&p);
        s
    });

    let cnts = count_occurrences(&result);

    let min = cnts.values().min().unwrap();
    let max = cnts.values().max().unwrap();

    max - min
}

fn main() {
    let input = include_str!("../../input-14.txt");

    let parsed = parse::parse(input).unwrap().1;

    println!("Part I:  {}", part1(&parsed));
    println!("Part II: {}", part2(&parsed));
}

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();
}

import std/[sequtils, strutils, sugar, os, tables, algorithm, math]

proc parseInput(inputFile: string): (string, Table[string, string]) =
  var input = collect(newSeq):
    for line in inputFile.lines: line

  result[0] = input[0]
  input = input[2 .. ^1]

  for line in input:
    let splitLine = line.split(" -> ")
    result[1][splitLine[0]] = splitLine[1]

func polymerize(polyTemplate: string, rules: Table[string, string]): string = 
  for idx in 0 .. polyTemplate.len - 2:
    result = result & 
             polyTemplate[idx] & 
             rules.getOrDefault(polyTemplate[idx .. idx + 1], "")
  result = result & polyTemplate[^1]

func makePolymer(polyTemplate: string, rules: Table[string, string], numIters: int): string =
  result = polyTemplate
  for iter in 1 .. numIters:
    result = polymerize(result, rules)

func analyzePolymer(polymer: string): int =
  var analysis = polymer.toCountTable()
  analysis.sort()
  result = analysis.values().toSeq()[0] - analysis.values().toSeq()[^1]

# I think this is a great joke
func toPairymer(polymer: string): CountTable[string] = 
  for idx in 0 .. polymer.len - 2:
    let pair = polymer[idx .. idx + 1]
    result[pair] = result.getOrDefault(pair, 0) + 1

func pairymerize(pairymer: CountTable[string], 
                 rules: Table[string, string]): CountTable[string] =
  for (pair, num) in pairs(pairymer):
    let ruleOutput = rules[pair]
    let pair1 = pair[0] & ruleOutput
    let pair2 = ruleOutput & pair[1]
    result[pair1] = result.getOrDefault(pair1, 0) + num
    result[pair2] = result.getOrDefault(pair2, 0) + num

func makePairymer(pairymer: CountTable[string], 
                  rules: Table[string, string], 
                  numIters: int): CountTable[string] =
  result = pairymer
  for iter in 1 .. numIters:
    result = pairymerize(result, rules)

func analyzePairymer(pairymer: CountTable[string]): int =
  var letterCounts: CountTable[char]
  for (pair, num) in pairs(pairymer):
    letterCounts[pair[0]] = letterCounts.getOrDefault(pair[0], 0) + num
    letterCounts[pair[1]] = letterCounts.getOrDefault(pair[1], 0) + num
  # This may be off by one but you can always guess lmao
  result = int(round((letterCounts.largest()[1] - letterCounts.smallest()[1]) / 2))

proc main(inputFile: string) =
  var (polymer, rules) = parseInput(inputFile) # Oh, so I can do that
  polymer = makePolymer(polymer, rules, 10)
  echo analyzePolymer(polymer)

  var pairymer = polymer.toPairymer() # Not wasting those iterations!
  pairymer = makePairymer(pairymer, rules, 30)
  echo pairymer.analyzePairymer()

when is_main_module:
  main(paramStr(1))

use crate::lib::aoc;
use std::collections::HashMap;

pub fn run() {
    let lines = aoc::get_lines("./inputs/day14.in");
    let template = lines[0].clone();
    let rules: HashMap<String, String> =
        lines[2..].iter().fold(HashMap::new(), |mut rules, line| {
            if let Some((pair, insertion)) = line.split_once(" -> ") {
                rules.insert(pair.to_owned(), insertion.to_owned());
            }
            rules
        });

    // Part 1
    let mut polymer = template;
    for _ in 0..10 {
        let mut polymer_chars = polymer.chars().peekable();
        let mut new_polymer = String::new();
        for i in 0..polymer.len() {
            let ith_char = polymer_chars.next().unwrap();
            new_polymer.push(ith_char);

            if i + 1 < polymer.len() {
                new_polymer.push_str(
                    rules
                        .get(&format!("{}{}", ith_char, polymer_chars.peek().unwrap()))
                        .expect("Missing some rule!"),
                );
            }
        }
        polymer = new_polymer;
    }

    let mut result_p1: Vec<u32> = polymer
        .chars()
        .fold(HashMap::<char, u32>::new(), |mut accum, element| {
            let count = accum.get(&element).unwrap_or(&0) + 1;
            accum.insert(element, count);

            accum
        })
        .values()
        .cloned()
        .collect();
    result_p1.sort_unstable();

    let result_p1 = result_p1.last().unwrap() - result_p1[0];

    aoc::output(14, "result", result_p1, 0);
}

use crate::lib::aoc;
use std::collections::HashMap;

fn do_steps(
    rules: &HashMap<String, char>,
    polymer: HashMap<(char, char), u64>,
    mut ends_with: (char, char),
) -> (HashMap<(char, char), u64>, (char, char)) {
    let mut new_polymer = HashMap::new();
    for (pair, count) in &polymer {
        let new_char = *rules
            .get(&format!("{}{}", pair.0, pair.1))
            .expect("Missing some rule!");

        // Now we can make two pairs
        let new_pair1 = (pair.0, new_char);
        let new_pair2 = (new_char, pair.1);
        new_polymer.insert(new_pair1, new_polymer.get(&new_pair1).unwrap_or(&0) + count);
        new_polymer.insert(new_pair2, new_polymer.get(&new_pair2).unwrap_or(&0) + count);

        if *pair == ends_with {
            ends_with = new_pair2;
        }
    }
    (new_polymer, ends_with)
}

fn get_result(polymer: &HashMap<(char, char), u64>, ends_with: &(char, char)) -> u64 {
    let mut count_vec: Vec<u64> = polymer
        .iter()
        .fold(
            HashMap::<char, u64>::new(),
            |mut accum, ((ith_char, jth_char), pair_count)| {
                let ith_char_count = accum.get(ith_char).unwrap_or(&0) + pair_count;
                let jth_char_count = accum.get(jth_char).unwrap_or(&0) + 1;

                if (*ith_char, *jth_char) == *ends_with {
                    accum.insert(*ith_char, ith_char_count);
                    accum.insert(*jth_char, jth_char_count);
                } else {
                    accum.insert(*ith_char, ith_char_count);
                }

                accum
            },
        )
        .values()
        .cloned()
        .collect();
    count_vec.sort_unstable();

    count_vec.last().unwrap() - count_vec[0]
}

pub fn run() {
    // Parse the input
    let lines = aoc::get_lines("./inputs/day14.in");
    let template = lines[0].clone();
    let rules: HashMap<String, char> = lines[2..].iter().fold(HashMap::new(), |mut rules, line| {
        if let Some((pair, insertion)) = line.split_once(" -> ") {
            rules.insert(
                pair.to_owned(),
                insertion.chars().next().expect("Some rule is invalid!"),
            );
        }
        rules
    });

    // Set up the polymer HashMap
    let mut template_chars = template.chars().peekable();
    let mut polymer: HashMap<(char, char), u64> = HashMap::new();
    let mut ends_with: (char, char) = ('0', '0');
    for i in 0..template.len() - 1 {
        let ith_char = template_chars.next().unwrap();
        let jth_char = *template_chars.peek().unwrap();
        let pair = (ith_char, jth_char);
        polymer.insert(pair, polymer.get(&pair).unwrap_or(&0) + 1);

        // We need to keep track of the ending pairs so that we can count chars correctly
        if i == template.len() - 2 {
            ends_with = pair;
        }
    }

    // ~~~ Part 1 ~~~
    for _ in 0..10 {
        let (new_polymer, new_ends_with) = do_steps(&rules, polymer.clone(), ends_with);
        polymer = new_polymer;
        ends_with = new_ends_with;
    }

    let result_p1 = get_result(&polymer, &ends_with);

    // ~~~ Part 2 ~~~
    for _ in 10..40 {
        let (new_polymer, new_ends_with) = do_steps(&rules, polymer.clone(), ends_with);
        polymer = new_polymer;
        ends_with = new_ends_with;
    }
    let result_p2 = get_result(&polymer, &ends_with);

    aoc::big_output(14, "result", result_p1, result_p2);
}

I could probably make it run a bit faster (not that it's particularly slow right) if I passed the number of steps to my do_steps function and iterated there so that I didn't have to pass the map to the function so much.

I guess I could also just keep track of the very last letter and add a count of 1 for that letter instead of keeping track of the last pair.

defmodule AdventOfCode.Solution.Year2021.Day14 do
  def part1(input) do
    measure_frequencies_at(input, 10)
  end

  def part2(input) do
    measure_frequencies_at(input, 40)
  end

  defp measure_frequencies_at(input, step) do
    {frequencies, rules, first_char, last_char} = parse_input(input)

    frequencies
    |> Stream.iterate(&insert(&1, rules))
    |> Enum.at(step)
    |> bigram_frequencies_to_unigram_frequencies(first_char, last_char)
    |> Map.values()
    |> Enum.min_max()
    |> then(fn {least_common_count, most_common_count} ->
      most_common_count - least_common_count
    end)
  end

  defp insert(freqs, rules) do
    freqs
    |> Enum.flat_map(fn {[char_a, char_b] = chars, count} ->
      rules[chars]
      |> then(fn insertion ->
        [
          {[char_a, insertion], count},
          {[insertion, char_b], count}
        ]
      end)
    end)
    |> Enum.group_by(&elem(&1, 0), &elem(&1, 1))
    |> Enum.map(fn {k, counts} -> {k, Enum.sum(counts)} end)
  end

  defp bigram_frequencies_to_unigram_frequencies(freqs, first_char, last_char) do
    freqs
    |> Enum.flat_map(fn {[char_a, char_b], freq} ->
      [
        {char_a, freq},
        {char_b, freq}
      ]
    end)
    |> Enum.group_by(&elem(&1, 0), &elem(&1, 1))
    |> Enum.map(fn {k, counts} -> {k, Enum.sum(counts)} end)
    |> Enum.into(%{}, fn
      {char, doubled_freq} -> {char, ceil(doubled_freq / 2)}
    end)
    |> then(fn freqs ->
      if first_char == last_char do
        Map.update!(freqs, first_char, &(&1 + 1))
      else
        freqs
      end
    end)
  end

  defp parse_input(input) do
    [first_char | rest_chars] = String.to_charlist(input)
    last_char = List.last(rest_chars)

    input
    |> String.split("\n\n", trim: true)
    |> then(fn [polymer, rules] ->
      {bigram_frequencies(polymer), parse_rules(rules), first_char, last_char}
    end)
  end

  defp bigram_frequencies(string) do
    string
    |> String.to_charlist()
    |> Enum.chunk_every(2, 1, :discard)
    |> Enum.frequencies()
  end

  defp parse_rules(rules) do
    rules
    |> String.split("\n", trim: true)
    |> Enum.map(fn <<char_a, char_b, " -> ", insertion>> -> {[char_a, char_b], insertion} end)
    |> Enum.into(%{})
  end
end