example.py 1.1 KB

1234567891011121314151617181920212223242526272829303132333435363738394041
  1. from itertools import chain
  2. class BinaryTreeNode(object):
  3. def __init__(self, value):
  4. self.value = value
  5. self.left = None
  6. self.right = None
  7. def set_left(self, node):
  8. self.left = node
  9. return self # return self so we can chain calls
  10. def set_right(self, node):
  11. self.right = node
  12. return self # return self so we can chain calls
  13. def __str__(self):
  14. return "<%d: %s, %s>" % (self.value, str(self.left), str(self.right))
  15. class IteratingBinaryTreeNode(BinaryTreeNode):
  16. def __iter__(self):
  17. left = self.left or []
  18. right = self.right or []
  19. return chain(iter(left), iter([self.value]), iter(right))
  20. if __name__ == "__main__":
  21. tree = BinaryTreeNode(1).set_left(BinaryTreeNode(0)).set_right(BinaryTreeNode(2))
  22. print tree
  23. iter_tree = (IteratingBinaryTreeNode(4) # Root
  24. .set_left(IteratingBinaryTreeNode(2) # L
  25. .set_left(IteratingBinaryTreeNode(1)) # L L
  26. .set_right(IteratingBinaryTreeNode(3))) # L R
  27. .set_right(IteratingBinaryTreeNode(5) # R L
  28. .set_right(IteratingBinaryTreeNode(6)))) # R R
  29. print list(iter(iter_tree))