Advent of Code 17.scala

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

enum CubeStatus {
  case inactive, active
}

case class Vector4(x: Int, y: Int, z: Int, w: Int) {
  def min(another: Vector4): Vector4 = {
    Vector4(x min another.x, y min another.y, z min another.z, w min another.w)
  }

  def max(another: Vector4): Vector4 = {
    Vector4(x max another.x, y max another.y, z max another.z, w max another.w)
  }

  def add(another: Vector4): Vector4 = {
    Vector4(x + another.x, y + another.y, z + another.z, w + another.w)
  }

  def multiply(another: Int): Vector4 = {
    Vector4(x * another, y * another, z * another, w * another)
  }
}

extension [A, B, C](m: Map[A, Map[B, C]])
  def adjust(k: A)(f: Map[B, C] => Map[B, C]) = m.updated(k, f(m.getOrElse(k, Map[B, C]())))

class CubeMap {
  var cubes: Map[Int, Map[Int, Map[Int, Map[Int, CubeStatus]]]] = Map()
  var min: Vector4 = Vector4(0, 0, 0, 0)
  var max: Vector4 = Vector4(0, 0, 0, 0)

  def set(pos: Vector4, status: CubeStatus): Unit = {
    cubes = cubes.adjust(pos.w)(_.adjust(pos.z)(_.adjust(pos.y)(_.updated(pos.x, status))))
    min = min.min(pos)
    max = max.max(pos)
  }

  def get(pos: Vector4): CubeStatus = {
    cubes.getOrElse(pos.w, Map()).getOrElse(pos.z, Map()).getOrElse(pos.y, Map()).getOrElse(pos.x, CubeStatus.inactive)
  }

  def copy(): CubeMap = {
    val result = CubeMap()
    result.cubes = cubes
    result.min = min
    result.max = max
    result
  }
}

@main def solve17() =
  def area(min: Vector4, max: Vector4): Seq[Vector4] = {
    for {
      newW <- min.w to max.w
      newZ <- min.z to max.z
      newY <- min.y to max.y
      newX <- min.x to max.x
    } yield Vector4(newX, newY, newZ, newW)
  }

  def nearBy(v: Vector4): Seq[Vector4] = {
    area(v.add(Vector4(-1, -1, -1, -1)), v.add(Vector4(1, 1, 1, 1))).filter(_ != v)
  }

  def solve(cubeMap: CubeMap, delta: Vector4, step: Int): Int = {
    if 0 == step then {
      area(cubeMap.min, cubeMap.max).map { v =>
        if cubeMap.get(v) == CubeStatus.active then 1 else 0
      }.sum
    } else
      val nextCubeMap = cubeMap.copy()
      area(cubeMap.min.add(delta.multiply(-1)), cubeMap.max.add(delta)).foreach { v =>
        val activeCount = nearBy(v).map { v => if cubeMap.get(v) == CubeStatus.active then 1 else 0 }.sum
        cubeMap.get(v) match {
          case CubeStatus.active =>
            if activeCount != 2 && activeCount != 3 then
              nextCubeMap.set(v, CubeStatus.inactive)
          case _ =>
            if activeCount == 3 then
              nextCubeMap.set(v, CubeStatus.active)
        }
      }
      solve(nextCubeMap, delta, step - 1)
  }

  def solve1(cubeMap: CubeMap, step: Int): Int = {
    solve(cubeMap, Vector4(1, 1, 1, 0), step)
  }

  def solve2(cubeMap: CubeMap, step: Int): Int = {
    solve(cubeMap, Vector4(1, 1, 1, 1), step)
  }

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

  val inputs = Iterator.continually(StdIn.readLine())
  val cubeMap = CubeMap()
  for {
    (ys, y) <- inputs.takeWhile(_ != null).toVector.zipWithIndex
    (xs, x) <- ys.toVector.zipWithIndex
  } yield cubeMap.set(Vector4(x, y, 0, 0), if xs == '#' then CubeStatus.active else CubeStatus.inactive)

  val answer1 = solve1(cubeMap, 6)
  println(answer1)

  val answer2 = solve2(cubeMap, 6)
  println(answer2)