python tree

Monday, July 21, 2008

I implement a simple tree with the DFS and BFS defined using python generators.

class Tree:
def __init__(self,data):
self.children = []
self.data = data

def dfs(self):
yield self
if len(self.children) == 0:
return
for child in self.children:
for c in child.dfs():
yield c

def bfs(self):
yield self
queue = list(self.children)
while len(queue) > 0:
first = queue[0]
queue = queue[1:]
queue.extend(first.children)
yield first


if __name__ == "__main__":
root = Tree(-1)
root.children = [Tree(i) for i in range(5)]
for i in range(5):
root.children[i].children = [Tree(i*j) for j in range(3)]

#-1,0,0,0,0...
for node in root.dfs():
print node.data

print
print

#-1,0,1,2,3,4,0,0,0,0,1,2,0,2,4...
for node in root.bfs():
print node.data

0 comments: