Advent of Code 11.scala

Advent of Code 11.scala

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

@main def solve11() =

  val endOfMap: Char = 'X'
  
  def seat(seats: Vector[Vector[Char]], x: Int, y: Int): Char = {
    seats.lift(y).flatMap(_.lift(x)).getOrElse(endOfMap)
  }

  def look(seats: Vector[Vector[Char]], x: Int, y: Int, dx: Int, dy: Int): Char = {
    val result = seat(seats, x + dx, y + dy)
    result match {
      case 'L' | '#' | `endOfMap` =>
        result
      case _ => 
        look(seats, x + dx, y + dy, dx, dy)
    }
  }
  
  val directions = Vector(
    (-1, -1), (-1, 0), (-1, 1), (0, -1), (0, 1), (1, -1), (1, 0), (1, 1)
  )

  def step_q1(seats: Vector[Vector[Char]]): Vector[Vector[Char]] = {
    val width = seats.head.size
    val height = seats.size
    (0 until height).toVector.map { y =>
      (0 until width).toVector.map { x =>
        seat(seats, x, y) match {
          case '.' =>
            '.'
          case 'L' =>
            if directions.exists((dx, dy) => seat(seats, x + dx, y + dy) == '#') then
              'L'
            else
              '#'
          case '#' =>
            if 4 <= directions.count((dx, dy) => seat(seats, x + dx, y + dy) == '#') then
              'L'
            else
              '#'
        }
      }
    }
  }

  def step_q2(seats: Vector[Vector[Char]]): Vector[Vector[Char]] = {
    val width = seats.head.size
    val height = seats.size
    (0 until height).toVector.map { y =>
      (0 until width).toVector.map { x =>
        seat(seats, x, y) match {
          case '.' =>
            '.'
          case 'L' =>
            if directions.exists((dx, dy) => look(seats, x, y, dx, dy) == '#') then
              'L'
            else
              '#'
          case '#' =>
            if 5 <= directions.count((dx, dy) => look(seats, x, y, dx, dy) == '#') then
              'L'
            else
              '#'
        }
      }
    }
  }

  def solve(seats: Vector[Vector[Char]], step: (Vector[Vector[Char]] => Vector[Vector[Char]])): Int = {
    val next = step(seats)
    if next == seats then
      seats.map(_.count(_ == '#')).sum
    else
      solve(next, step)
  }

  val in = new FileInputStream("example11-1.in")
  System.setIn(in)

  val inputs = Iterator.continually(StdIn.readLine()).takeWhile(_ != null).toVector
  val seats = inputs.map { line => line.toVector }
  val answer1 = solve(seats, step_q1)
  println(answer1)
  val answer2 = solve(seats, step_q2)
  println(answer2)