Method and helper method snippet from a quadtree node class. Returns square 2d array representation of the quadtree.

quadtree_to_array.py

# Just came out of the oven. Unclean and fugly.
# 
# Key observation is that, when walking the array left-to-right
# and top-to-bottom, the "nw"/"ne"/"sw"/"se" attributes for the
# chain of subnodes down to the bottom have simple, power-of-2
# toggle periods.
#
# For example, to access the top-leftmost cell in a 32x32
# quadtree called foo:
#
# foo.nw.nw.nw.nw.nw
#
# The next cell to the right is:
#
# foo.nw.nw.nw.nw.ne
#
# That rightmost "w"/"e" has a period of 1. Here are the toggle
# periods for the rest:
#
#     512 16  256 8  128 4   64 2   32 1
# foo. N  W .  N  W . N  W . N  W . N  W
#
# As you can see, there are math.log(32,2)=5 levels. "e"/"w" periods
# increase from right to left by powers of two, and then "n"/"s" wraps
# around again.
#

def to_array(self):
    raw_array = self._to_array_helper()
    size = int(math.sqrt(len(raw_array)))
    result = []
    while len(raw_array) >= size:
        result.append(raw_array[:size])
        raw_array = raw_array[size:]
    return result

def _to_array_helper(self):
    class Periodic_Coord():
        def __init__(self,letter_pair,period):
            self._letters = letter_pair
            self._value = 0
            self._period = period
            self._counter = 0
        def get(self):
            current = self._letters[self._value]
            self._counter = (self._counter + 1) % self._period
            if self._counter == 0:
                self._value = self._value ^ 1
            return current
    result = []
    level = 1
    probe = self
    while type(probe.nw) != int:
        probe = probe.nw
        level += 1
    horiz_periods = [2**x for x in xrange(level)]
    vert_periods = [2**x for x in xrange(level,level*2)]
    horiz_coords = map(lambda p: Periodic_Coord(("w","e"),p), horiz_periods)
    vert_coords = map(lambda p: Periodic_Coord(("n","s"),p), vert_periods)
    horiz_coords.reverse()
    vert_coords.reverse()
    def coord(lev):
        h = horiz_coords[lev].get()
        v = vert_coords[lev].get()
        return v+h
    i = 0
    iMax = 2**(2*level)
    while i <= iMax:
        cell = self
        L = 0
        while type(cell) != int:
            cell = getattr(cell,coord(L))
            L = (L + 1) % len(horiz_coords)
        result.append(cell)  
        i += 1
    return result