Day22.scala

package day22

import scalaz._, Scalaz._

object Day22 extends App {
  sealed trait Equipment

  object Equipment {
    case object Climbing extends Equipment
    case object Torch    extends Equipment
    case object Neither  extends Equipment
  }

  sealed abstract class RiskLevel(val risk: Long,
                                  char: Char,
                                  val compatibleEquipment: Set[Equipment]) {
    override def toString: String = char.toString
  }

  object RiskLevel {
    def fromErosion(erosionLevel: Long): RiskLevel =
      erosionLevel % 3 match {
        case 0 => Rocky
        case 1 => Wet
        case 2 => Narrow
      }

    case object Rocky  extends RiskLevel(0, '.', Set(Equipment.Climbing, Equipment.Torch))
    case object Wet    extends RiskLevel(1, '=', Set(Equipment.Climbing, Equipment.Neither))
    case object Narrow extends RiskLevel(2, '|', Set(Equipment.Torch, Equipment.Neither))
  }

  final case class Point(x: Int, y: Int) {
    def +(other: Point): Point           = Point(x + other.x, y + other.y)
    def manhattanDist(other: Point): Int = (x - other.x).abs + (y - other.y).abs
  }
  object Point {
    val Origin = Point(0, 0)
  }

  final case class Cave(depth: Long) {
    private lazy val erosionLevel: Point => Long =
      Memo.immutableHashMapMemo[Point, Long] { point =>
        (geoIndex(point) + depth) % 20183
      }

    private def geoIndex(point: Point): Long = point match {
      case Point(0, 0) => 0L
      case Point(0, y) => y * 48271L
      case Point(x, 0) => x * 16807L
      case Point(x, y) =>
        erosionLevel(Point(x - 1, y)) * erosionLevel(Point(x, y - 1))
    }

    def riskLevel(point: Point): RiskLevel =
      RiskLevel.fromErosion(erosionLevel(point))

    def totalRisk(target: Point): Long =
      (for {
        x     <- 0 to target.x
        y     <- 0 to target.y
        point = Point(x, y)
        if point != Point.Origin && point != target
      } yield riskLevel(point).risk).sum

    def toMap(target: Point, pos: Option[Point] = None): String = {
      def toMapLine(y: Int) =
        List
          .tabulate(target.x) { x =>
            val point = Point(x, y)
            if (pos.contains(point)) "X"
            else riskLevel(point).toString
          }
          .mkString

      List
        .tabulate(target.y)(y => toMapLine(y))
        .mkString("", "\n", "\n")
    }

    def minTimeTo(target: Point): Int = {
      final case class Node(pos: Point, risk: RiskLevel, equipment: Equipment) {
        def distanceTo(other: Node): Int =
          pos.manhattanDist(other.pos) + (if (equipment == other.equipment) 0 else 7)

        def optimisticDistanceTo(other: Node): Int = pos.manhattanDist(other.pos)

        def neighbors: Set[Node] = reequip ++ moves

        private def moves: Set[Node] =
          for {
            delta     <- Set(Point(0, 1), Point(1, 0), Point(0, -1), Point(-1, 0))
            targetPos = pos + delta
            if targetPos.x >= 0 && targetPos.y >= 0
            targetRisk = if (targetPos == target) RiskLevel.Rocky else riskLevel(targetPos)
            if targetRisk.compatibleEquipment.contains(equipment)
          } yield Node(targetPos, targetRisk, equipment)

        private def reequip: Set[Node] =
          (risk.compatibleEquipment - equipment).map(targetEquipment =>
            copy(equipment = targetEquipment))
      }

      val initialNode = Node(Point.Origin, RiskLevel.Rocky, Equipment.Torch)
      val targetNode  = Node(target, RiskLevel.Rocky, Equipment.Torch)
      val searcher = new AStar[Node](distance = _.distanceTo(_).toDouble,
                                     neighbors = _.neighbors,
                                     heuristic = _.optimisticDistanceTo(targetNode).toDouble)

      val Some(path) = searcher.search(initialNode, isGoal = _ == targetNode)

      path.foreach { node =>
        println(s"Equipment: ${node.equipment}")
        println(toMap(target, pos = Some(node.pos)))
        println()
      }

      path.sliding(size = 2).toList.foldMap {
        case List(prev, next) =>
          prev.distanceTo(next)
      }
    }
  }

  val Target = Point(14, 778)
  val cave   = Cave(depth = 11541L)
  println(cave.totalRisk(Target))
  println(cave.minTimeTo(Target)) // Not 1089
}

Day22Test.scala

package day22

import day22.Day22.{Cave, Point, RiskLevel}
import org.scalatest._

final class Day22Test extends FlatSpec with Matchers {
  val cave = Cave(depth = 510L)

  "Day 22" should "compute risk levels" in {
    cave.riskLevel(Point(x = 1, y = 0)) shouldBe RiskLevel.Wet
  }

  it should "print a map" in {
    cave.toMap(target = Point(x = 16, y = 8)) shouldBe """.=.|=.|.|=.|=|=.
                                                         |.|=|=|||..|.=...
                                                         |.==|....||=..|==
                                                         |=.|....|.==.|==.
                                                         |=|..==...=.|==..
                                                         |=||.=.=||=|=..|=
                                                         ||.=.===|||..=..|
                                                         ||..==||=.|==|===
                                                         |""".stripMargin
  }

  it should "find the shortest path" in {
    cave.minTimeTo(Point(10, 10)) shouldBe 45
  }
}