Advent of Code 2021-9.scala

import java.io.FileInputStream
import scala.annotation.tailrec
import scala.io.StdIn

case class HeightMap(heightMap: Vector[Vector[Int]]) {
  def width(): Int = {
    heightMap.head.size
  }

  def height(): Int = {
    heightMap.size
  }

  def value(point: Point): Int = {
    heightMap.lift(point.y)
      .getOrElse(Vector())
      .lift(point.x)
      .getOrElse(10)
  }
}

case class Point(x: Int, y: Int) {
  def adjacentPoints(): Vector[Point] =
    Vector(Point(x - 1, y), Point(x + 1, y), Point(x, y - 1), Point(x, y + 1))
}

case class Basin(points: Vector[Point])

def isLowPoint(heightMap: HeightMap, point: Point): Boolean = {
  val value = heightMap.value(point)
  point.adjacentPoints().forall(p => value < heightMap.value(p))
}

def locationOfLowPoints(heightMap: HeightMap): Vector[Point] =
  (0 until heightMap.height()).flatMap { y =>
    (0 until heightMap.width()).flatMap { x =>
      if isLowPoint(heightMap, Point(x, y)) then
        Some(Point(x, y))
      else
        None
    }
  }.toVector

def solve9_1(inputs: Vector[Vector[Int]]): Int =
  val heightMap = HeightMap(inputs)
  locationOfLowPoints(heightMap).map { point => heightMap.value(point) + 1 }.sum

def findBasins(heightMap: HeightMap, lowPoints: Vector[Point]): Vector[Basin] =
  @tailrec
  def _findBasins(acc: Set[Point] = Set(), nextPoints: Vector[Point]): Basin =
    nextPoints.headOption match
      case Some(point) =>
        val value = heightMap.value(point)
        val newNextPoints = point.adjacentPoints().filter(p => value < heightMap.value(p) && 9 > heightMap.value(p))
        _findBasins(acc + point, nextPoints.tail ++ newNextPoints)
      case None =>
        Basin(acc.toVector)

  lowPoints.map(point => _findBasins(nextPoints = Vector(point)))


def solve9_2(inputs: Vector[Vector[Int]]): Int =
  val heightMap = HeightMap(inputs)
  val lowPoints = locationOfLowPoints(heightMap)
  val basins = findBasins(heightMap, lowPoints)
  basins.map(_.points.size).sortBy(-_).take(3).product

@main def solve9(): Unit =
  val in = new FileInputStream("example9-2.in")
  System.setIn(in)
  val inputs = Iterator.continually(StdIn.readLine())
    .takeWhile(line => null != line && line.trim.nonEmpty)
    .map(_.map(_.toInt - '0').toVector)
    .toVector
  println(solve9_1(inputs))
  println(solve9_2(inputs))