I am studying the following binary-search-tree implementation1 in python:
class TreeNode():
def __init__(self, key, val,
lc = None, rc = None, parent = None):
self.key = key
self.val = val
self.leftChild = lc
self.rightChild = rc
self.parent = parent
def hasLeftChild(self):
return not self.leftChild is None
def hasRightChild(self):
return not self.rightChild is None
# in-order iterator
def __iter__(self):
if self:
if self.hasLeftChild():
for elem in self.leftChild:
yield elem
yield self.key
if self.hasRightChild():
for elem in self.rightChild:
yield elem
class BinarySearchTree():
def __init__(self):
self.root = None
self.size = 0
def length(self):
return self.size
def __len__(self):
return self.size
def put(self, key, val):
if not self.root:
self.root = TreeNode(key, val)
else:
self._put(key, val, self.root)
def _put(self, key, val, currentNode):
if currentNode.key >= key:
if currentNode.hasLeftChild():
self._put(key, val, currentNode.leftChild)
else:
currentNode.leftChild = TreeNode(key, val, parent = currentNode)
else:
if currentNode.hasRightChild():
self._put(key, val, currentNode.rightChild)
else:
currentNode.rightChild = TreeNode(key, val, parent = currentNode)
Let's consider this example caller code:
from random import randint
bst = BinarySearchTree()
keys = [5, 30, 2, 40, 25, 4]
for key in keys:
bst.put(key, randint(1000, 2000))
# in-order iteration
for key in bst.root:
print(key)
producing the expected output:
2
4
5
25
30
40
What is happening behind the scenes when the for
loop used for the in-order iteration in the caller code is executed?
I know that the __iter__
method of the TreeNode
class recursively calls itself in the two for
loops for the left and right sub-trees.
What I don't manage to grasp is the effect of yield elem
in those for
loops:
if self.hasLeftChild(): for elem in self.leftChild: yield elem
and
if self.hasRightChild(): for elem in self.rightChild: yield elem
with respect to the effect of yield self.key
, which apparently concerns the usual implementation of the iterator protocol as explained in the docs.
More specifically, what is the content of the elem
variable that is returned by yield elem
in the above-mentioned nested for
loops, as opposed to the content of self.key
(which is the actual key data in the TreeNode
object) returned by yield self.key
?
- This is an excerpt from the implementation given in Problem Solving with Algorithms and Data Structures using Python by Brad Miller and David Ranum.