Aoc2021-day14.scala

package lascala.aoc2021.day11_20.day14

opaque type Polymer = Seq[Char]
opaque type PairInsertion = Map[(Char, Char), Char]

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

object Polymer:
  def apply(s: String): Polymer = s.toSeq

extension (p: Polymer)
  def step(pi: PairInsertion): Polymer =
    p.sliding(2).flatMap {
      case Seq(a, b) =>
        if pi contains (a, b) then Seq(a, pi((a, b))) else Seq(a)
    }.toSeq :+ p.last
  def toPairCount: PairCount =
    p.sliding(2).map {
      case Seq(a, b) => (a, b) -> BigInt(1)
    }.toMap
  def elementCount: Map[Char, Int] =
    p.groupBy(identity).view.mapValues(_.size).toMap
  def mkString: String = p.mkString

extension (pc: PairCount)
  def step(pi: PairInsertion): PairCount =
    val pc1 = for
      ((a, b), count) <- pc.toSeq
      ((a1, b1), count1) <- {
        if pi contains (a, b) then
          Seq((a, pi((a, b))) -> count, (pi((a, b)), b) -> count)
        else
          Seq((a, b) -> count)
      }
    yield
      ((a1, b1) -> count1)
    pc1.groupMapReduce(_._1)(_._2)(_ + _)
  def elementCount(last: Char): Map[Char, BigInt] =
    val firstCounts = pc.toSeq.map{ case ((a, b), count) => (a, count) } :+ (last -> BigInt(1))
    firstCounts.groupMapReduce(_._1)(_._2)(_ + _)


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 solve1(s: String): BigInt =
  val (polymer, pairInsertion) = parse(s)
  val last: Char = polymer.mkString.last
  val p1: PairCount = (0 until 10).foldLeft(polymer.toPairCount){
    case (pc, _) => pc.step(pairInsertion)
  }
  val elementCount = p1.elementCount(polymer.mkString.last)
  val counts = elementCount.values
  counts.max - counts.min
end solve1

def solve2(s: String): BigInt =
  val (polymer, pairInsertion) = parse(s)
  val last: Char = polymer.mkString.last
  val p1: PairCount = (0 until 40).foldLeft(polymer.toPairCount){
    case (pc, _) => pc.step(pairInsertion)
  }
  val elementCount = p1.elementCount(polymer.mkString.last)
  val counts = elementCount.values
  counts.max - counts.min
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")
  }