Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

python iterator through tree with list of children

i don't get the full grasp on python iterators, i got an object with a list of children, and i want to iterate through this structure. I want to get the same behaviour as with the printall function but with an iterator.

    class t:
            def __init__(self, i):
                    self.l = []
                    self.a = 0
                    for ii in range(i):
                            self.a = ii
                            self.l.append(t(i-1))

            def __iter__(self):
                    return self

            def next(self):
                    for i in self.l:
                            yield i.__iter__()
                    yield self

            def printall(self):
                    for i in self.l:
                            i.printall()
                    print self.a

hope thats enough information, thanks

edit:

i just want to be able to iterate through all the leafs of the tree and do something with the object, i.e. when i have an instance

    bla = t(3) 

i want to be able to go through every node with

    for x in bla:
            print x.a

for example. i want to be able to something with each x, i just have to access every child once

like image 570
Xtroce Avatar asked Aug 02 '11 15:08

Xtroce


1 Answers

It sounds like you want the iterator to act as a tree traversal. Study the itertools module and you can really go places.

from itertools import chain, imap

class t:
  def __init__(self, value):
    self.value = value
    self.children = []
  def __iter__(self):
    "implement the iterator protocol"
    for v in chain(*imap(iter, self.children)):
      yield v
    yield self.value

root = t(0)
root.children.append(t(1))
root.children.append(t(2))
root.children[1].children.append(t(3))
print list(iter(root))   # -> [1, 3, 2, 0]
print list(iter(root.children[1]))  # -> [3, 2]

EDIT: Below is the originally accepted implementation. It has a performance problem; I would remove it, but it seems wrong to remove content that was an accepted answer. It will fully traverse the entire structure, creating O(N*log[M](N)) generator objects (for a balanced tree with branching factor M containing N total elements), before yielding any values. But it does produce the desired result with a simple expression.

(The above implementation visits areas of the tree on demand and has only O(M+log[M](N)) generator objects in memory at a time. In both implementations, only O(log[M](N)) levels of nested generators are expected.)

from itertools import chain

def isingle(item):
  "iterator that yields only a single value then stops, for chaining"
  yield item

class t:
  # copy __init__ from above
  def __iter__(self):
    "implement the iterator protocol"
    return chain(*(map(iter, self.children) + [isingle(self.value)]))
like image 63
wberry Avatar answered Oct 10 '22 06:10

wberry