Trees
class Node(object): # ...
def cousins(self):
"""Find nodes on the same level as this node."""
# START SOLUTION
# Find our depth (0=root node, 1=child-of-root, etc)
# and the root of the tree. We'll do this by walking up the tree
# until we find the root, counting our steps up.
current = self
sought_depth = 0
while current.parent is not None:
sought_depth += 1
current = current.parent
# Now, "current" is the root node, as we navigated to the top
#
# Now, let's navigate back down to that level. We'll do this
# recursively, but we could do it with a stack, too.
# Set of cousins we'll find.
#
# Sets are mutable, so we'll just mutate this copy in place as
# we find cousins.
cousins = set()
# Recursive function to find examine a node, decide whether to
# consider it a cousin, and to explore its children.
#
# We could make this a free-standing function outside of the class,
# of a method of the class -- but we don't need access to the
# `self`, and by nesting this, we'll have access to the
# `cousins` and `sought_depth` variables, so we conveniently don't
# need to pass them around explicitly.
def _cousins(node, curr_depth):
# Base case: we're a leaf node, so can't look further
if node is None:
return
# Second base case: we're at the right level, so don't need
# to look further
if curr_depth == sought_depth:
cousins.add(node)
return
for child in node.children:
# Look at the children, noting that we're one level deeper then.
_cousins(child, curr_depth + 1)
_cousins(node=current, curr_depth=0)
# We don't want to include the original node
cousins.remove(self)
return cousins