Trials, tribulations, and tributaries (Rainfall challenge)

This is a fairly simple Ruby solution to the August 2016 community challenge.

Ok, the title is way too dramatic. But tributary is a nice word, and the alliteration was too good to pass up.

Besides, konijn started it!

I didn't bother to output a 2D "map" of the basins, I just print their sizes as the challenge requires. (Printing a map wouldn't be a difficult thing graft onto this, but I just wanted to keep things clean.)

Anyway, the script:

# Simple class that models a point on the terrain
class Cell
  attr_reader :altitude

  # Init with the cell's altitude
  def initialize(altitude)
    @altitude = altitude
    @neighbors = []
  end

  # Add another cell instance as a neighbor of the receiver
  def add_neighbor(cell)
    @neighbors << cell
  end

  # Returns the "drain" of the receiver (lowest of lower neighbors),
  # or nil if no such neighbor exists
  def drain
    @drain ||= @neighbors
      .select { |neighbor| neighbor.altitude < altitude }
      .min_by(&:altitude)
  end

  # Returns the sink of the receiver, i.e. the drain's drain (ad fundum),
  # or self if the cell is itself a sink
  def sink
    @sink ||= drain.nil? ? self : drain.sink
  end
end

# Helper method. Given an array of cells, register adjacent
# cells as neighbors of each other
def link_cells(cells)
  cells.each_cons(2) do |a, b|
    a.add_neighbor(b)
    b.add_neighbor(a)
  end
end

# Get the terrain size (first line of input)
size = gets.chomp.to_i

# Init the altitude cells (as a 2D array) from input
terrain = size.times.map do
  gets.strip.split.map { |altitude| Cell.new(altitude.to_i) }
end

# Register neighbors (first row-wise, then column-wise)
terrain.each { |row| link_cells(row) }
terrain.transpose.each { |col| link_cells(col) }

# Divide into basins (cells sharing same sink) and print basin sizes
basins = terrain.flatten.group_by(&:sink).values
puts basins.map(&:size).sort.reverse.join(" ")

It passes the test cases given in the challenge. It also works with non-square

I'm self-reviewing here, because I found some rough spots. Do not let this keep you, dear reader, from writing a review of your own!

For this review I'll leave the overall approach as-is, and just focus on the details of its implementation. I'm nit-picking like crazy, but that's the fun part.

As mentioned in the question, I'm not super happy with the 2-step "initialization" needed to get a useful Cell instance. Mostly because it makes the Cell class tricky to use. Call #drain or #sink without adding any neighbors, and you get useless results. But at least you can add the neighbors later and try again. But call either method with some but not all neighbors added, it'll memoize possibly incorrect result. Adding neighbors after that won't correct the results.

One solution would be to wrap the Cell usage in an "outer" API. For instance a Terrain class or module, with a #basin_sizes(input) (or initializer) method that uses Cell internally to produce the answer. There'd still be a 2-step "initialization", but it'd be hidden from the outside world and (presumably) handled correctly.

Another (partial) solution would be to null the memoized results whenever a new neighbor is added. This way there wouldn't be the same risk of getting stale results, if, say, #drain was called before all the neighbors were added. It wouldn't solve the problem of #sink being memoized elsewhere, though: A cell's drain might change, but if other cells have already memoized a sink that relies on the cell's previous drain, they won't automatically update. Water doesn't flow uphill.

Speaking of memoizing, there's some inefficiency in the way it's used here. For a sink-cell, #drain will of course evaluate to nil. But memoizing a nil value with ||= has no effect: Next time the method's called, it'll evaluate the right-hand expression all over again - and again and again, always producing nil. The point of memoizing should of course be that the right-hand expression isn't run again.

One way to solve this could be to use a 2nd instance variable to track whether the results been calculated yet:

def drain
  return @drain if @drain_found
  @drain_found = true
  @drain = ...
end

But that's not very elegant.

Alternatively, one could initialize the memoization variable to a separate value/type, like false, and check for that:

def drain
  return @drain if @drain != false
  @drain = ...
end

However, all of this is actually academic. Since the method that's actually called when solving is #sink, and memoization works well for #sink (no risk of a nil value). So in the end, the memoization of #drain doesn't actually matter: It's only called once anyway. In fact, #drain might as well be private.

But perhaps a better alternative overall would be to move the "lowest lower cell"-selection from #drain to #add_neighbor:

def add_neighbor(cell)
  reference = drain ? drain.altitude : altitude # get lowest altitude
  @drain = cell if cell.altitude < reference    # update drain if lower
end

Edit: One could cut that down to

@drain = cell if cell.altitude < (drain || self).altitude

but it's probably clearer as-is.

Obviously, add_neighbor might not be the best name for this method anymore. It's more like "set drain to this cell, maybe, if you feel like it"... yeah, I don't know exactly what the name should be. Maybe just #register_neighbor or something.

Anyway, it'd shift the drain identification to be "up front", and the #drain method could be replaced with an attr_reader.

Implementing the above yields a smaller Cell class, which is nice:

class Cell
  attr_reader :altitude, :drain

  def initialize(altitude)
    @altitude = altitude
  end

  def register_neighbor(cell)
    reference = drain ? drain.altitude : altitude
    @drain = cell if cell.altitude < reference
  end

  def sink
    @sink ||= drain ? drain.sink : self
  end
end

There's still the issue of #sink being memoized, so it shouldn't be called before all neighbors have been added/tested. As far as I can tell, this is a "lesser of two evils" question. Having the memoization there short-circuits the pathfinding, so the cell#sink → cell#sink → cell#sink → ... call stack doesn't have to go full-depth every time, which I feel is worth it. Again, the entire thing can be hidden behind a single API method to protect from premature memoization.

But I won't bother with that. The code's not intended for use outside of this script anyway, so for all intents and purposes the script itself is the wrapper.

StackExchange Code Review Q#136772, answer score: 2