watershed_final.py

import string

from collections import defaultdict
from heapq import heappop, heappush

class UnionFind(object):
    def __init__(self, size):
        # not enough
        self.size = defaultdict(int)
        self.weights = range(size)

    def union(self, a, b):
        # Get index roots
        root_a = self._root(a)
        root_b = self._root(b)

        # Added some balancing
        if self.size[root_a] > self.size[root_b]:
            self.weights[root_b] = self.weights[root_a]
            self.size[root_a] += self.size[root_b]
        else:
            self.weights[root_a] = self.weights[root_b]
            self.size[root_b] += self.size[root_a]

    def iter_weights(self):
        for idx in xrange(len(self.weights)):
            yield self._root(idx)

    def _root(self, idx):
        last = idx
        while self.weights[last] != last:
            last = self.weights[last]
        return last


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

    res = []
    # Added a second sort parameter (order_of_importance) to maintain
    # order between N, W, E, S
    order_of_importance = iter(xrange(0, 4))

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

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

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

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


def main():
    with open('B-large-practice.in') as f:
        n_cases = int(f.readline())

        for n 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):
                try:
                    cost, _, neighbour_id = get_neighbours(matrix, idx, width, height)
                except IndexError:
                    continue
                if cost < el:
                    uf.union(idx, neighbour_id)

            print "Case #%s:" % (n + 1)
            letters = iter(string.ascii_lowercase)
            d = {}
            iter_weights = uf.iter_weights()
            for _ in xrange(height):
                cs = []
                for _ in xrange(width):
                    num = next(iter_weights)
                    try:
                        cs.append(d[num])
                    except KeyError:
                        char = d[num] = next(letters)
                        cs.append(char)
                print ' '.join(cs)


if __name__ == '__main__':
    main()