Aoc2021-day14.scala

package lascala.aoc2021.day11_20.day14

opaque type Polymer = Seq[Char]
object Polymer:
  def apply(s: String): Polymer = s.toSeq

opaque type PairInsertion = Map[(Char, Char), Char]
opaque type PairCount     = Map[(Char, Char), BigInt]

extension [K](s: Seq[(K, BigInt)])
  def mapReduce: Map[K, BigInt] = s.groupMapReduce(_._1)(_._2)(_ + _)

extension (p: Polymer)
  def toPairCount: PairCount = p
    .sliding(2)
    .toSeq
    .groupMapReduce { case Seq(a, b) => (a, b) }(_ => BigInt(1))(_ + _)
  def mkString: String = p.mkString

extension (pc: PairCount)
  def step(pi: PairInsertion): PairCount = pc.toSeq.flatMap {
    case ((a, b), count) =>
      pi.get((a, b)).fold(Seq((a, b) -> count)) { c =>
        Seq((a, c) -> count, (c, b) -> count)
      }
  }.mapReduce

  def elementCountWithLast(last: Char): Map[Char, BigInt] = pc.toSeq
    .map { case ((a, _), count) => (a, count) }
    .appended(last -> BigInt(1))
    .mapReduce

def parse(s: String): (Polymer, PairInsertion) =
  val Array(template, rules) = s.split("\n\n")
  val pairInsertion: PairInsertion = rules
    .split("\n")
    .map { line =>
      val Array(k, v) = line.split(" -> ")
      (k.head, k.last) -> v.head
    }
    .toMap

  (Polymer(template), pairInsertion)

def solve(s: String, numberOfStep: Int): BigInt =
  val (polymer, pairInsertion) = parse(s)
  val last: Char               = polymer.mkString.last
  val p1: PairCount = (0 until numberOfStep)
    .foldLeft(polymer.toPairCount) { (pc, _) => pc.step(pairInsertion) }
  val elementCount = p1.elementCountWithLast(last)
  val counts       = elementCount.values
  counts.max - counts.min

def solve1(s: String): BigInt =
  solve(s, 10)
end solve1

def solve2(s: String): BigInt =
  solve(s, 40)
end solve2

@main def part1: Unit =
  val ans = solve1(input)
  println(ans)

@main def part2: Unit =
  val ans = solve2(input)
  println(ans)

//val input = """KOKHCCHNKKFHBKVVHNPN

Aoc2021-day14Test.scala

package lascala.aoc2021.day11_20.day14

import minitest.SimpleTestSuite
import hedgehog.minitest.HedgehogSupport
import hedgehog.*

object Day14Test extends SimpleTestSuite with HedgehogSupport:

  val testInput = """NNCB

CH -> B
HH -> N
CB -> H
NH -> C
HB -> C
HC -> B
HN -> C
NN -> C
BH -> H
NC -> B
NB -> B
BN -> B
BB -> N
BC -> B
CC -> N
CN -> C"""

  example("day14 - solve 1") {
    solve1(testInput) ==== 1588
  }

  example("day14 - solve 2") {
    solve2(testInput) ==== BigInt("2188189693529")
  }