Advent of Code 2021 Day 22 Neal Wu solution

aoc-2021-day22.scala

// using scala 3.0.2

package day22


import scala.util.Using
import scala.io.Source

import scala.collection.mutable

import scala.util.chaining.*

import Command.*

@main def part1(): Unit =
  println(s"The solution is ${part1(readInput())}")

@main def part2(): Unit =
  println(s"The solution is ${part2(readInput())}")

def readInput(): String =
  Using.resource(Source.fromFile("day22"))(_.mkString)

case class Dimension(min: Int, max: Int):
  require(min <= max)

  def isSubset(d: Dimension): Boolean =
    min >= d.min && max <= d.max

extension (x1: Int)
  infix def by (x2: Int): Dimension = Dimension(x1, x2)

case class Cuboid(xs: Dimension, ys: Dimension, zs: Dimension)

enum Command:
  case On, Off

case class Step(command: Command, cuboid: Cuboid)

object Indexed:
  opaque type Index = Int

  object Index:

    def apply(x: Int, y: Int, z: Int): Index =
      val x0 = x & 0x3ff
      val y0 = y & 0x3ff
      val z0 = z & 0x3ff
      (x0 << 20) | (y0 << 10) | z0

    def decode(i: Int): Index = i

    extension (i: Index)
      def x: Int =
        ((i >> 20) & 0x3ff)
      def y: Int =
        ((i >> 10) & 0x3ff)
      def z: Int =
        (i & 0x3ff)
      def encoded: Int =
        i

import Indexed.*

def run(steps: IndexedSeq[Step]): Long = {

  val X0, Y0, Z0 = mutable.ArrayBuffer.empty[Int]

  for Step(_, Cuboid(xs, ys, zs)) <- steps do
    X0 += xs.min
    X0 += xs.max + 1
    Y0 += ys.min
    Y0 += ys.max + 1
    Z0 += zs.min
    Z0 += zs.max + 1

  val X = IArray.from(X0).sorted
  val Y = IArray.from(Y0).sorted
  val Z = IArray.from(Z0).sorted

  val XL = X.zipWithIndex.toMap
  val YL = Y.zipWithIndex.toMap
  val ZL = Z.zipWithIndex.toMap

  val N = X.size

  val g = mutable.BitSet.empty

  for Step(tpe, Cuboid(xs, ys, zs)) <- steps do
    val x0 = XL(xs.min)
    val x1 = XL(xs.max + 1)
    val y0 = YL(ys.min)
    val y1 = YL(ys.max + 1)
    val z0 = ZL(zs.min)
    val z1 = ZL(zs.max + 1)
    val value = tpe == Command.On
    for
      x <- x0 until x1
      y <- y0 until y1
      z <- z0 until z1
    do
      val i = Index(x, y, z).encoded
      if value then g += i else g -= i

  val sum =
    class Sum extends (Int => Unit):
      private var acc = 0L

      def add(i: Index): Unit =
        add(i.x, i.y, i.z)

      def add(x: Int, y: Int, z: Int): Unit =
        acc += 1L * (X(x + 1) - X(x)) * (Y(y + 1) - Y(y)) * (Z(z + 1) - Z(z))

      def apply(enc: Int): Unit = add(Index.decode(enc))

      def run = this.tap(g.foreach).acc

    Sum().run

  sum
}

def isInit(cuboid: Cuboid): Boolean =
  Seq(cuboid.xs, cuboid.ys, cuboid.zs).forall(_.isSubset(-50 by 50))

type Parser[A] = PartialFunction[String, A]

val NumOf: Parser[Int] =
  case s if s.matches(raw"-?\d+") => s.toInt

val DimensionOf: Parser[Dimension] =
  case s"${NumOf(begin)}..${NumOf(end)}" => begin by end

val CuboidOf: Parser[Cuboid] =
  case s"x=${DimensionOf(xs)},y=${DimensionOf(ys)},z=${DimensionOf(zs)}" => Cuboid(xs, ys, zs)

val CommandOf: Parser[Command] =
  case "on" => On
  case "off" => Off

val StepOf: Parser[Step] =
  case s"${CommandOf(command)} ${CuboidOf(cuboid)}" => Step(command, cuboid)

def part1(input: String) =
  run(input.linesIterator.map(StepOf).filter(s => isInit(s.cuboid)).toIndexedSeq)

def part2(input: String) =
  run(input.linesIterator.map(StepOf).toIndexedSeq)