Day16.scala

package sungkmi.aoc2020.day16

type Ranges = Map[String, Seq[Range]]

extension (rs: Ranges)
  def containsNumber(n: Int): Boolean = rs.values.flatten.exists(_ `contains` n)

def parse(s: String): (Ranges, Seq[Int], Seq[Seq[Int]]) =
  val Array(input1, input2, input3) = s `split` "\n\n"
  val ranges: Ranges = input1.split("\n").map:
    line =>
      val Array(key, value) = line `split` ": "
      val valueRanges = value.split(" or ").toSeq.map:
        r =>
          val Array(s, e) = r.split("-").map(_.toInt)
          s to e
      key -> valueRanges
  .toMap
  val yourTicket = input2.split("\n").tail.head.split(",").map(_.toInt).toSeq
  val nearbyTickets = input3.split("\n").tail.map(_.split(",").map(_.toInt).toSeq).toSeq
  (ranges, yourTicket, nearbyTickets)

def solve1(input: String): Int =
  val (rs, _, nearbyTickets) = parse(input)
  nearbyTickets.flatten.filterNot(rs.containsNumber).sum

def solve2(input: String): Map[String, Int] = 
  val (rs, yourTicket, nearbyTickets) = parse(input)
  val ticketKeySets: Seq[Seq[(Int, Set[String])]] = nearbyTickets.map:
    nearbyTicket =>
      nearbyTicket.map:
        number =>
          val keys = for {
            (k, v) <- rs if v.exists(_ `contains` number)
          } yield k
          (number -> keys.toSet)
  .filterNot(_.exists(_._2.isEmpty))
  val candidates = ticketKeySets.foldLeft(Vector.fill(rs.size)(rs.keySet)):
    (keySets, ticketKeySet) =>
      keySets.zip(ticketKeySet).map:
        (set1, set2) =>
          set1 `intersect` set2._2
  val candidateMap = candidates.zipWithIndex.map:
    (set, index) => index -> set
  .toMap

  @annotation.tailrec
  def loop(
    candidateMap: Map[Int, Set[String]],
    ans: Map[Int, String]
  ): (Map[Int, Set[String]], Map[Int, String]) =
//    println(s"===> $candidateMap, $ans")
    val confirmed = for{
      (k, v) <- candidateMap if v.size == 1
    } yield (k, v.head)
//    println(s"======> confirmed: $confirmed")
    if confirmed.isEmpty then (candidateMap, ans) else
      val newAns = ans ++ confirmed
      val newCandidates = (candidateMap -- confirmed.keySet).view.mapValues:
        set => set -- confirmed.values.toSet
      .toMap
      loop(newCandidates, newAns)
  val ansMap = loop(candidateMap, Map.empty)._2
  yourTicket.zipWithIndex.map:
    (n, i) => ansMap(i) -> n
  .toMap

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

@main def part2: Unit = 
  val ansMap = solve2(input)
  val ans = (for {
    (k, v) <- ansMap if k `startsWith` "departure"
  } yield BigInt(v)).product
  println(ans)

lazy val input: String = """departure location: 35-796 or 811-953

Day16Test.scala

package sungkmi.aoc2020.day16

class Day16Test extends munit.FunSuite {

  val testInput = """class: 1-3 or 5-7
row: 6-11 or 33-44
seat: 13-40 or 45-50

your ticket:
7,1,14

nearby tickets:
7,3,47
40,4,50
55,2,20
38,6,12"""

  test("parse") {
    val expectedRanges: Ranges = Map(
      "class" -> Seq(1 to 3, 5 to 7),
      "row" -> Seq(6 to 11, 33 to 44),
      "seat" -> Seq(13 to 40, 45 to 50),
    )
    val expectedYourTicket = Seq(7, 1, 14)
    val expectedNearbyTickets = Seq(
      Seq(7, 3, 47),
      Seq(40, 4, 50),
      Seq(55, 2, 20),
      Seq(38, 6, 12),
    )

    val (rangeSet, yourTicket, nearbyTickets) = parse(testInput)

    assertEquals(rangeSet, expectedRanges)
    assertEquals(yourTicket, expectedYourTicket)
    assertEquals(nearbyTickets, expectedNearbyTickets)

  }

  test("day16 solve1") {  
    assertEquals(solve1(testInput), 71)
  }

  test("day16 solve2") {  

    val testInput2 = """class: 0-1 or 4-19
row: 0-5 or 8-19
seat: 0-13 or 16-19

your ticket:
11,12,13

nearby tickets:
3,9,18
15,1,5
5,14,9"""

    val expected = Map("class" -> 12, "row" -> 11, "seat" -> 13)
    assertEquals(solve2(testInput2), expected)
  }
}