Rainfall challenge - DFS

rainfall_challenge_dfs.py

# coding=utf-8
"""
Rainfall challenge using DFS
https://codereview.stackexchange.com/questions/38500/rainfall-challenge

O(n²) time and space
"""
import sys
import collections


def rainfall_challenge_dfs(grid):
    visited = set()
    sinks = {}
    n = len(grid)
    delta = [(-1, 0), (1, 0), (0, -1), (0, 1)]

    def dfs(row, col):
        this_coord = (row, col)
        visited.add((row, col))
        min_elev = sys.maxint
        min_neighbor = None

        for dy, dx in delta:
            nrow = row + dy
            ncol = col + dx

            if 0 <= nrow < n and 0 <= ncol < n and grid[nrow][ncol] < min_elev:
                min_elev = grid[nrow][ncol]
                min_neighbor = (nrow, ncol)

        if min_elev > grid[row][col]:
            sink = this_coord
        elif min_neighbor in sinks:
            sink = sinks[min_neighbor]
        else:
            sink = dfs(*min_neighbor)

        sinks[this_coord] = sink
        return sink

    for i in xrange(n):
        for j in xrange(n):
            if (i, j) not in visited:
                dfs(i, j)

    counts = collections.defaultdict(lambda: 0)

    for i in xrange(n):
        for j in xrange(n):
            counts[sinks[(i, j)]] += 1

    return sorted([v for v in counts.itervalues()], reverse=True)

rainfall_challenge_uf.py

# coding=utf-8
"""
Rainfall challenge using Union-Find
https://codereview.stackexchange.com/questions/38500/rainfall-challenge

O(n²) time and space
"""
import sys
import collections


def rainfall_challenge_uf(grid):
    s = len(grid)
    n = s * s
    parent = [i for i in xrange(n)]
    rank = [0] * n

    def index(row, col):
        return row * s + col

    def find(p):
        while p != parent[p]:
            parent[p] = parent[parent[p]]
            p = parent[p]
        return p

    def union(p, q):
        root_p = find(p)
        root_q = find(q)

        if root_p == root_q:
            return

        if rank[root_p] < rank[root_q]:
            parent[root_p] = root_q
        elif rank[root_p] > rank[root_q]:
            parent[root_q] = root_p
        else:
            parent[root_q] = root_p
            rank[root_p] += 1

    deltas = [(-1, 0), (1, 0), (0, -1), (0, 1)]

    for row in xrange(s):
        for col in xrange(s):
            min_elev = sys.maxint
            min_nbor = None

            for dx, dy in deltas:
                nrow = row + dy
                ncol = col + dx

                if 0 <= nrow < s and 0 <= ncol < s and grid[nrow][ncol] < min_elev:
                    min_elev = grid[nrow][ncol]
                    min_nbor = (nrow, ncol)

            if min_elev < grid[row][col]:
                union(index(row, col), index(*min_nbor))

    counts = collections.defaultdict(lambda: 0)

    for row in xrange(s):
        for col in xrange(s):
            counts[find(index(row, col))] += 1

    return sorted([c for c in counts.itervalues()], reverse=True)