23. Code Jam
By: Steve Krenzel
Previous Index Next

This challenge by far involved the most code.

The problem is: You're given a topology and need to determine how water will flow on it. When you calculate where each point of water will flow, you also determine a basin where the water would build up. Each point in the topology is in a basin, your job is to find which basin and give that basin a label.

The topology is given in a grid of M-by-N integers. Starting at each integer you can just traverse to the neighbor that is lowest until there are no neighbors that are lower, at which point you've found the bottom of that basin. This is a brute force approach but it's a pretty simple recursive function and works more than fast enough for even the larger data set.

The code I'm presenting here is pretty simple and short at only 30 lines long. In my humble opinion this is a better reference than the 45 lines long solution given by Google in their analysis. I'm not saying it's a better reference because it's shorter, but simply because it's a little more straight foward and easier to wrap your head around. Here is my code:

from sys import stdin

def neighbors(topo, x, y):
    if y > 0:                yield (y - 1, x)
    if x > 0:                yield (y, x - 1)
    if x < len(topo[0]) - 1: yield (y, x + 1)
    if y < len(topo) - 1   : yield (y + 1, x)

def find_basin(topo, x, y):
    xs, ys = x, y
    for i, j in neighbors(topo, x, y):
        if topo[i][j] < topo[ys][xs]:
            xs, ys = j, i
    return (x, y) if (x, y) == (xs, ys) else find_basin(topo, xs, ys)

def print_basins(topo):
    labels = {}
    for y in xrange(len(topo)):
        for x in xrange(len(topo[y])):
            basin = find_basin(topo, x, y)
            print labels.setdefault(basin, chr(ord('a') + len(labels))),

# Function to read in a line from stdin and split it into ints
read_line = lambda: map(int, stdin.readline().split())

for m in xrange(int(stdin.readline())):
    print "Case #%d:"%(m+1)
    h, w = read_line()
    print_basins([read_line() for i in range(h)])

To try it just save the code in a file (e.g. watersheds.py), download sample input from the challenge site, and run "python watersheds.py < B-small-practice.in" at the command line.

About the author
I'm Steve Krenzel, a software engineer and co-founder of Thinkfuse. Contact me at steve@thinkfuse.com.
Previous Next