Day9b.rb

require 'set'

class Cell
  def initialize(cave, x, y)
    @cave = cave
    @x = x
    @y = y
  end

  attr_reader :x, :y

  def height
    @cave.height_at(x, y)
  end

  def neighbors
    @cave.neighbors_at(x, y)
  end

  def basin
    @cave.basin_at(x, y)
  end

  def low?
    height < neighbors.map(&:height).min
  end
end

class Cave
  def initialize(map)
    @map = map
    @height = map.length
    @width = map[0].length
    @basin_by_coords = {}
    @basins = []
  end

  attr_reader :width, :height, :basins

  def [](x, y)
    Cell.new(self, x, y)
  end

  def height_at(x, y)
    @map[y][x]
  end

  def neighbors_at(x, y)
    ns = []
    ns << self[x - 1, y] if x > 0
    ns << self[x + 1, y] if x + 1 < width
    ns << self[x, y - 1] if y > 0
    ns << self[x, y + 1] if y + 1 < height
    ns
  end

  def basin_at(x, y)
    basin = @basin_by_coords[[x, y]]
    if !basin
      cell = self[x, y]
      return if cell.height >= 9
      if cell.low?
        basin = Set.new
        @basins << basin
      else
        basin = cell.neighbors.min_by(&:height).basin or return
      end
      basin.add [x, y]
      @basin_by_coords[[x, y]] = basin
    end
    basin
  end
end


map = $<.map {|l| l.chomp.chars.map(&:to_i)}

cave = Cave.new(map)

risk = 0
(0...cave.height).each do |y|
  (0...cave.width).each do |x|
    cave[x, y].basin
  end
end

print cave.basins.map(&:count).sort {|x, y| y <=> x}.take(3).reduce(:*)