uf.py

import string

class UnionFind(object):
    def __init__(self, size):
        # not enough
        alpha = iter(string.ascii_lowercase)
        self.weights = [alpha.next() for _ in xrange(size)]

    def union(self, a, b):
        # Add top-left rule
        try:
            smaller_letter, bigger_letter = sorted((
                self.weights[a], self.weights[b]
            ))
        except IndexError:
            import ipdb; ipdb.set_trace()

        # Make all bigger == smaller
        for idx in xrange(len(self.weights)):
            if self.weights[idx] == bigger_letter:
                self.weights[idx] = smaller_letter


def get_neighbours(matrix, idx, width, height):
    y = idx / width
    x = idx % width

    res = []
    order_of_importance = iter(xrange(0, 4))

    if y - 1 >= 0:
        key = (y-1) * width + x
        res.append(
            (matrix[key], next(order_of_importance), key)
        )

    if x - 1 >= 0:
        key = y * width + (x - 1)
        res.append(
            (matrix[key], next(order_of_importance), key)
        )

    if x + 1 < width:
        key = y * width + (x + 1)
        res.append(
            (matrix[key], next(order_of_importance), key)
        )

    if y + 1 < height:
        key = (y+1) * width + x
        res.append(
            (matrix[key], next(order_of_importance), key)
        )
    
        # Better use a heap
    return sorted(res)[0]


def main():
    with open('watershed.txt') as f:
        n_cases = int(f.readline())

        for _ in xrange(n_cases):
            # Pluck width and height
            height, width = map(int, f.readline().split(' '))

            # Create union find datastructure
            uf = UnionFind(width * height)

            matrix = []
            for _ in xrange(height):
                matrix.extend(map(int, f.readline().split(' ')))

            for idx, el in enumerate(matrix):
                cost, _, neighbour_id = get_neighbours(matrix, idx, width, height)
                if cost < el:
                    uf.union(idx, neighbour_id)

            iter_weights = iter(uf.weights)
            for _ in xrange(height):
                print ' '.join(next(iter_weights) for _ in xrange(width))


if __name__ == '__main__':
    main()