Advent of Code 2021-12.scala

import java.io.FileInputStream
import scala.annotation.tailrec
import scala.io.StdIn

case class Cave(name: String):
  val isLarge: Boolean = name.exists(_.isUpper)

case class Graph(inputs: Vector[(Cave, Cave)]):
  val links: Map[Cave, Vector[Cave]] = inputs.foldLeft(Map[Cave, Vector[Cave]]()) { (acc, v) =>
    acc.updated(v._1, acc.getOrElse(v._1, Vector()) :+ v._2)
      .updated(v._2, acc.getOrElse(v._2, Vector()) :+ v._1)
  }

trait Path:
  val last: Cave
  def move(cave: Cave): Option[Path]

case class Q1Path(last: Cave, limited: Set[Cave]) extends Path:
  def move(cave: Cave): Option[Path] =
    if limited.contains(cave) then
      None
    else
      Some(Q1Path(cave, if cave.isLarge then limited else limited + cave))
object Q1Path:
  def apply(cave: Cave): Path =
    Q1Path(cave, if cave.isLarge then Set() else Set(cave))

case class Q2Path(last: Cave, limited: Set[Cave], isTwiceVisited: Boolean) extends Path:
  def move(cave: Cave): Option[Path] =
    if cave.name == "start" || (limited.contains(cave) && isTwiceVisited) then
      None
    else
      Some(Q2Path(cave, if cave.isLarge then limited else limited + cave, isTwiceVisited || limited.contains(cave)))
object Q2Path:
  def apply(cave: Cave): Path =
    Q2Path(cave, if cave.isLarge then Set() else Set(cave), false)


@tailrec
def traverse(graph: Graph, stack: List[Path], acc: Int): Int =
  stack match
    case head :: tail =>
      val nexts = graph.links.getOrElse(head.last, Vector())
      val updatedPaths = nexts.flatMap(head.move).toList
      val (end, moving) = updatedPaths.partition(_.last.name == "end")
      traverse(graph, moving ::: tail, acc + end.size)
    case _ =>
      acc

def solve12_1(inputs: Vector[(Cave, Cave)]): Int =
  val graph = Graph(inputs)
  traverse(graph, List(Q1Path(Cave("start"))), 0)

def solve12_2(inputs: Vector[(Cave, Cave)]): Int =
  val graph = Graph(inputs)
  traverse(graph, List(Q2Path(Cave("start"))), 0)


@main def solve12(): Unit =
  val in = new FileInputStream("example11-2.in")
  System.setIn(in)
  val inputs = Iterator.continually(StdIn.readLine())
    .takeWhile(line => null != line && line.trim.nonEmpty)
    .map(line => line.split("-"))
    .map { case Array(from, to) => (Cave(from), Cave(to)) }
    .toVector
  println(solve12_1(inputs))
  println(solve12_2(inputs))