Generating random maze using disjoint sets

mazegen.py

import random
import sys
import numpy
 
X = 20
Y = 20
 
# the following creats a two dimensional array
# of indexes. Each cell now has a number, and we can convert co-ordinates
# into numbers via this array
indexes = numpy.arange(X*Y).reshape(X,Y)
# These 1d arrays correspond to the data previously inside the cell class
right_walls = numpy.ones(X*Y, bool)
down_walls = numpy.ones(X*Y, bool)
 
 
 
def gen_walls():
    # create an array which holds information about walls
    # There will be three columns
    # First column: parent id
    # Second collumn: neighbor id
    # Third column: True if direction of wall is to the right
    total_right_walls = (X - 1) * Y
    total_down_walls =  (Y - 1) * X
    total_wall_count = total_right_walls + total_down_walls
    walls = numpy.empty( (total_wall_count, 3), int)
 
    # fill up the walls array with right walls
    walls[:total_right_walls,0] = indexes[:-1,:].flatten()
    walls[:total_right_walls,1] = indexes[1:,:].flatten()
    walls[:total_right_walls,2] = True
 
    # fill up the walls array with down walls
    walls[total_right_walls:,0] = indexes[:,:-1].flatten()
    walls[total_right_walls:,1] = indexes[:,1:].flatten()
    walls[total_right_walls:,2] = False
 
    return walls
 
 
class UnionFind(object):
    def __init__(self, sections):
        self.parent = numpy.arange(sections)
        self.rank = numpy.ones(sections)
        self.final = numpy.ones(sections, bool)
 
    def find(self, a):
        # in python, local variable access is much faster then attributes on classes
        # when speed is really really neccesary, its worthwhile to assign any attributes
        # accessed multiple times to local variabled
        parents = self.parent
        parent = parents[a]
        if not self.final[parent]:
            parent = self.find(parent)
            parents[a] = parent
        return parent
 
    def union(self, a, b):
        find = self.find
        a_parent = find(a)
        b_parent = find(b)
        if a_parent != b_parent:
            rank = self.rank
            parent = self.parent
            final = self.final
            a_rank = rank[a_parent]
            b_rank = rank[b_parent]
            if a_rank < b_rank:
                parent[a_parent] = b_parent
                final[a_parent] = False
            elif b_rank < a_rank:
                parent[b_parent] = a_parent
                final[b_parent] = False
            else:
                parent[b_parent] = a_parent
                final[b_parent] = False
                rank[a_parent] += 1
            return True
        else:
            return False
 
 
 
def break_walls():
    walls = gen_walls()
    numpy.random.shuffle(walls)
 
    # Smash some walls
    uf = UnionFind( X*Y )
 
    for parent, neighbour, right in walls:
        # If the cells on either side of the wall aren't already connected,
        # destroy the wall
        if uf.union(parent, neighbour):
            if right:
                right_walls[parent] = False
            else:
                down_walls[parent] = False
 
 
def output(blocked):
    sys.stdout.write(" X"[blocked])
   
def draw_maze():
    # Draw the maze
    for x in range(2*X+1):
        output(True)
    sys.stdout.write('\n')
 
    for y in xrange(Y):
        output(True)
        for x in xrange(X):
            output(False)
            output(right_walls[indexes[x,y]])
        sys.stdout.write('\n')
        output(True)
        for x in xrange(X):
            output(down_walls[indexes[x,y]])
            output(True)
        sys.stdout.write('\n')
 
def main():
    break_walls()
    draw_maze()
       
if __name__ == '__main__':
    main()