basins.py

"""
Palantir interview question.

The goal is to, given an elevation map of some land, print out a
list of basin sizes in descending order, where a basin is a set of
cells of land that all drain to a common point. One cell drains to
another if among the first cell and its neighbors (i.e. cells adjacent,
horizonatlly, vertically or diagonally) the second cell has the lowest
elevation.

Thus, given the input:

3
1 5 2
2 4 7
3 6 9

The output should be:

7 2

Since there are two basins:

A A B
A A B
A A A
"""
from collections import defaultdict
import sys


class Cell(object):
    def __init__(self, x, y, elevation):
        self.x = x
        self.y = y
        self.elevation = elevation
        self.target = self
        self._basin = None

    def union(self, other):
        if self.target.elevation > other.elevation:
            self.target = other

    def union_many(self, others):
        for other in others:
            self.union(other)

    def get_neighbors(self, cells):
        """
        Get the neighbors of this cell.
        """
        for i in [-1, 0, 1]:
            for j in [-1, 0, 1]:
                try:
                    yield cells[self.x + i, self.y + j]
                except KeyError:
                    pass

    @property
    def basin(self):
        if self._basin:
            return self._basin
        if self.target == self:
            self._basin = self
            return self
        else:
            self._basin = self.target.basin
            return self.target.basin

    def __repr__(self):
        return "Cell(%s, %s, %s)" % (self.x, self.y, self.elevation)


def altitudes_to_cells(height, altitudes):
    """
    Convert a two dimensional array of altitudes into a list of cells.

    >>> height = 3
    >>> altitudes = [[1, 5, 2], [2, 4, 7], [3, 6, 9]]
    >>> altitudes_to_cells(height, altitudes)
    [Cell(0, 1, 5), Cell(1, 2, 7), Cell(0, 0, 1), Cell(2, 1, 6), Cell(1, 1, 4), Cell(2, 0, 3), Cell(2, 2, 9), Cell(1, 0, 2), Cell(0, 2, 2)]
    """
    cells = {}
    for x in range(height):
        for y in range(height):
            altitude = altitudes[x][y]
            cells[x, y] = Cell(x, y, altitude)
    for cell in cells.values():
        cell.union_many(cell.get_neighbors(cells))
    return cells.values()


def group_cells_by_basin(cells):
    """
    Given a list of cells, group them by common basins. Returns a map from
    basin to set of cells that have that drain into that basin.
    """
    basins = defaultdict(set)
    for cell in cells:
        basins[cell.basin].add(cell)
    return dict(basins)


def get_basins(cells):
    """
    Get a list of basins.

    >>> height = 3
    >>> altitudes = [[1, 5, 2], [2, 4, 7], [3, 6, 9]]
    >>> cells = altitudes_to_cells(height, altitudes)
    >>> len(get_basins(cells))
    2
    """
    return group_cells_by_basin(cells).values()


def parse_elevation_map(lines):
    """
    Parse an elevation map input. Returns the height of the map and a
    two-dimensional array of altitudes.

    Input should begin with a line with one integer, S, the height (and width)
    of the map. The next S lines will each contain a row of the map, each with
    S integers - the elevations of the S cells in the row.  Some farmers have
    small land plots such as the examples below, while some have larger plots.

    - Assumes the elevation maps are square.

    - Assumes that in no case will the total size of the farmland exceed
      1000x1000.

    """
    height = int(lines.readline())
    altitudes = []
    for line in lines:
        altitudes.append([int(i) for i in line.split()])
    return height, altitudes


def main():
    """
    Print a list of a space-separated list of the basin sizes, in descending
    order.
    """
    height, altitudes = parse_elevation_map(sys.stdin)
    cells = altitudes_to_cells(height, altitudes)
    basins = get_basins(cells)
    print " ".join((str(i) for i in sorted([len(basin) for basin in basins], reverse=True)))

if __name__ == "__main__":
    main()