Advent of Code 2021-14.scala

import java.io.FileInputStream
import scala.io.StdIn

def groupTuples[A, B](seq: Seq[(A, B)]): Map[A, Seq[B]] = seq.groupBy(_._1).view.mapValues(_.map(_._2)).toMap

def step(pairs: Map[(Char, Char), BigInt], pairInsertions: Map[(Char, Char), Char]): Map[(Char, Char), BigInt] =
  val newPairs = pairs.toVector.flatMap(pair => {
    val ((c0, c1), count) = pair
    val newChar = pairInsertions((c0, c1))
    Vector((c0, newChar) -> count, (newChar, c1) -> count)
  })
  groupTuples(newPairs).view.mapValues(_.sum).toMap

def answer(polymerTemplate: String, pairInsertions: Map[(Char, Char), Char], loop: Int): BigInt =
  val pairs: Map[(Char, Char), BigInt] = polymerTemplate.sliding(2).map(s => (s(0), s(1))).toVector.groupBy(x => x).view.mapValues(v => BigInt(v.size)).toMap
  val loopedPairs: Map[(Char, Char), BigInt] = (0 until loop).foldLeft(pairs)((acc, idx) => step(acc, pairInsertions))
  val result = groupTuples(loopedPairs.toVector.flatMap { case ((k0, k1), v) => Vector(k0 -> v, k1 -> v) } :+
    (polymerTemplate.head -> BigInt(1)) :+
    (polymerTemplate.last -> BigInt(1))
  ).view.mapValues(_.sum).toVector.sortBy(_._2)
  (result.last._2 - result.head._2) / 2


def solve14_1(polymerTemplate: String, pairInsertions: Map[(Char, Char), Char]): BigInt =
  answer(polymerTemplate, pairInsertions, 10)

def solve14_2(polymerTemplate: String, pairInsertions: Map[(Char, Char), Char]): BigInt =
  answer(polymerTemplate, pairInsertions, 40)

@main def solve14(): Unit =
  val in = new FileInputStream("example14-2.in")
  System.setIn(in)
  val polymerTemplate = StdIn.readLine()
  StdIn.readLine()
  val inputs = Iterator.continually(StdIn.readLine())
    .takeWhile(line => null != line && line.trim.nonEmpty)
    .map(line => ((line(0), line(1)) -> line(6)))
    .toVector
  println(solve14_1(polymerTemplate, inputs.toMap))
  println(solve14_2(polymerTemplate, inputs.toMap))