Aoc2021-day9.scala

package lascala.aoc2021.day1_10.day9

case class Point(row: Int, column: Int)

case class HeightMap(
  map: Map[Point, Int],
  numberOfRows: Int,
  numberOfColumns: Int,
)

extension (hm: HeightMap)
  def height(point: Point): Option[Int] = hm.map.get(point)

  def lowPointsHeight: Seq[Int] = for
    row <- 0 until hm.numberOfRows
    column <- 0 until hm.numberOfColumns
    point = Point(row, column)
    Some(h) = height(point)
    if h < point.neighbors.flatMap(height).min
  yield h

  def basins: Seq[Int] =

    val points: Seq[Point] = for
      row <- 0 until hm.numberOfRows
      column <- 0 until hm.numberOfColumns
      if height(Point(row, column)).get < 9
    yield Point(row, column)

    val flow: Map[Point, Point] = {
      for
        p <- points
      yield 
        val lowest = (p :: p.neighbors)
          .flatMap{ n => height(n).map( h => (h, n)) }
          .minBy(_._1)
          ._2
        (p -> lowest)
    }.toMap

    flow.toSeq.sortBy{ case (f, t) => (f.row, f.column)} foreach println

    val basins0: Set[List[Point]] = points.map(List(_)).toSet

    val lookup0: Map[Point, List[Point]] =
      points.map(p => (p -> List(p))).toMap
    
    val (basins, lookup) = flow.foldLeft((basins0, lookup0)){
      case ((basins, lookup), (from, to)) =>

        val fromBasin = lookup(from)
        val toBasin = lookup(to)
        val combined = fromBasin ::: toBasin

        val basins1 = basins - fromBasin - toBasin + combined
        val lookup1 = lookup + (from -> combined) + (to -> combined)
//        println("=======")
//        println(s"$from -> $to")
//        println("---")
//        for
//          row <- 0 until hm.numberOfRows
//        yield
//          println((0 until hm.numberOfColumns).map{
//            column =>
//              lookup1.get(Point(row, column)).fold(Point(row, column))(_.minBy(p => height(p).get))
//          }.mkString(" "))
//        println("---")
//        basins1 foreach println
//        println("---")
//        lookup1 foreach println
//        println("=======")
        (basins1, lookup1)      
    }


    basins.toSeq.map(_.size)
  end basins


end extension

extension (p: Point)
  def neighbors: List[Point] =
    List(
      Point(p.row - 1, p.column),
      Point(p.row + 1, p.column),
      Point(p.row, p.column - 1),
      Point(p.row, p.column + 1),
    )
end extension

object HeightMap:
  def parse(s: String): HeightMap =
    val lines = s.split("\n")
    val numberOfRows = lines.length
    val numberOfColumns = lines.head.length
    val map = for
      row <- 0 until numberOfRows
      column <- 0 until numberOfColumns
      point = Point(row, column)
      h = lines(row)(column).toString.toInt
    yield point -> h.toInt
    HeightMap(map.toMap, numberOfRows, numberOfColumns)
  end parse

def solve1(s: String): BigInt =
  HeightMap.parse(s).lowPointsHeight.map(_ + 1).sum
end solve1

def solve2(s: String): BigInt =
  val basins = HeightMap.parse(s).basins.sorted.reverse
  println(basins)
  basins.take(3).map(BigInt(_)).product
end solve2

@main def part1: Unit =
  val ans = solve1(input)
  println(ans)

@main def part2: Unit =
  val ans = solve2(input)
  println(ans)

//val input = 

Aoc2021-day9Test.scala

package lascala.aoc2021.day1_10.day9

import minitest.SimpleTestSuite
import hedgehog.minitest.HedgehogSupport
import hedgehog.*

object Day9Test extends SimpleTestSuite with HedgehogSupport:

  val testInput = """2199943210
3987894921
9856789892
8767896789
9899965678"""

  example("day9 - solve1") {
    solve1(testInput) ==== 15
  }
  
  example("day9 - solve2") {
    solve2(testInput) ==== 1134
  }

end Day9Test