Maze generation in Python

When I was in school, I created (as an assignment) a program that created mazes and allowed the user to run through it in 3D. I used a recursive backtracking algorithm to create the maze. While this isn't a bad way to do it, I wanted to try implementing a different solution in python.

Enter Kruskal. Kruskal's algorithm is an algorithm used to find a minimum spanning tree for a graph. If this doesn't sound like a good way to build a maze, just hang on. I'm getting there. A lot of the time, the most difficult part of applying an algorithm to a problem is fitting it into some sort of data structure. A maze makes a perfect graph--every square is a node in the graph. Every square that is not separated by a wall is connected.

For more details on graph theory or Kruskal's algorithm, Wikipedia is your friend.

Back to the point: I love how simple the python implementation is.

width, height = 20, 20

# create a list of all walls 
# (all connections between squares in the maze)
# add all of the vertical walls into the list
walls = [(x,y,x+1,y)
    for x in range(width-1)
    for y in range(height)]
# add all of the horizontal walls into the list
walls.extend([(x,y,x,y+1)
    for x in range(width)
    for y in range(height-1)])

# create a set for each square in the maze
cell_sets = [set([(x,y)])
    for x in range(width)
    for y in range(height)]

# in Kruskal's algorithm, the walls need to be
# visited in order of weight
# since we want a random maze, we will shuffle 
# it and pretend that they are sorted by weight
walls_copy = walls[:]
random.shuffle(walls_copy)

for wall in walls_copy:
    set_a = None
    set_b = None

    # find the sets that contain the squares
    # that are connected by the wall
    for s in cell_sets:
        if (wall[0], wall[1]) in s:
            set_a = s
        if (wall[2], wall[3]) in s:
            set_b = s

    # if the two squares are in separate sets,
    # then combine the sets and remove the 
    # wall that connected them
    if set_a is not set_b:
        cell_sets.remove(set_a)
        cell_sets.remove(set_b)
        cell_sets.append(set_a.union(set_b))
        walls.remove(wall)

I'm sure there are things here that could be improved, but I really like the way Python works. If I were to implement this in C++, it might run faster (not that it needs to, the mazes are generated fast enough in Python) but it would take a lot longer to write and would end up with a lot more code.