Maze Generator (based on Kruskal Algorithm)

maze.py

# _ _ _ _ _ _ _ _ _ _
#|_   _____   _      |
#| |_|  _______|_| |_|
#|  _| |_  |  ___   _|
#|_  | | |_  | |_____|
#|_  | | |_  | | |   |
#|_    |_     _  | |_|
#|___|_  | |_|  ___  |
#|   |_  |  _|_|  _| |
#|_|_  | |___  |     |
#|___________|_|_|_|_|

import sys
from random import choice

N, S, E, W = 1, 2, 4, 8
DX = { str(E):1, str(W):-1, str(N): 0, str(S):0 }
DY = { str(E):0, str(W): 0, str(N):-1, str(S):1 }
OPPOSITE = { str(E):W, str(W):E, str(N):S, str(S):N}

class TreeNode(object):
    def __init__(self, parent=None, group=0):
        self.parent = parent
        self.group = group

    def is_root(self):
        return self.parent == None

    def get_parent(self):
        node = self
        while (not node.is_root()):
            node = node.parent
        return node

    def set_parent(self, parent):
        if (parent.is_root()):
            self.parent = parent
        else:
            self.parent = parent.get_parent()


class Tree(object):
    def __init__(self, root):
        self.root = root

    def get_root(self):
        return self.root.get_parent()

    def joinWith(self, otherTree):
        otherTree.get_root().set_parent(self.root)

class Maze(object):
    def __init__(self, width, height):
        self.sets = [[Tree(TreeNode(None, y*width + x)) for x in range(width)] for y in range(height)]
        self.grid = [[0 for x in range(width)] for y in range(height)]
        self.width = width
        self.height = height
        self._edges = []

        for y in range(height):
            for x in range(width):
                if y > 0:
                    self._edges.append([x,y,str(N)]) 
                if x > 0:
                    self._edges.append([x,y,str(W)])

        while len(self._edges) != 0:
            edge = choice(self._edges)
            self._edges.remove(edge)

            x, y, direction = edge
            nx, ny = x + DX[direction], y + DY[direction]

            set1, set2 = self.sets[y][x], self.sets[ny][nx]

            if set1.get_root() != set2.get_root():
                set1.joinWith(set2)
                self.grid[y][x] = self.grid[y][x] | int(direction)
                self.grid[ny][nx] = self.grid[ny][nx] | OPPOSITE[direction]
        
    def display(self):
        print " _" * (len(self.grid[0]) )
        for y in range(self.height):
            sys.stdout.write("|")
            for x in range(self.width):
                cell = self.grid[y][x]
                if cell == 0:
                    sys.stdout.write("\e[47m")
                if (cell & S != 0):
                    sys.stdout.write(" ")
                else:
                    sys.stdout.write("_")
                if (cell & E != 0):
                    if ((cell | self.grid[y][x+1]) & S):
                        sys.stdout.write(" ")
                    else:
                        sys.stdout.write("_")
                else:
                    sys.stdout.write("|")
                if cell == 0:
                    print "\e[m"
              
            sys.stdout.write("\n")


def main():
    maze = Maze(10,10);
    maze.display()

    print maze.grid

if __name__ == "__main__":
    main()