Skip to content

Instantly share code, notes, and snippets.

@Ceasar Ceasar/basins.py
Created Jan 13, 2014

Embed
What would you like to do?
"""
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()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.