Skip to content

Instantly share code, notes, and snippets.

@PirosB3
Created February 10, 2015 16:00
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save PirosB3/0abb75f393e381478d9b to your computer and use it in GitHub Desktop.
Save PirosB3/0abb75f393e381478d9b to your computer and use it in GitHub Desktop.
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 = []
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()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment