Tree Cousins

Whiteboard: Harder
Challenge: Medium
Concepts: Trees
Download: tree-cousins.zip
Solution: Tree Cousins: Solution

Visualize a tree like:

Nodes k and l are at same level (“cousins”). Similarly e, f, g, h, i, and j are cousins, as are b, c, and d.

In this problem, you’ll need to write a method for finding the nodes at the same level as a given node. It should return a set of those nodes at the same level, but not including the node it was called on.

For example:

>>> b.cousins() == {c, d}
True

>>> c.cousins() == {b, d}
True

>>> e.cousins() == {f, g, h, i, j}
True

>>> k.cousins() == {l}
True

The root node has no cousins:

>>> a.cousins() == set()   # empty set
True

We’ve given you a file, cousin.py, with a Node class with basic methods needed to create nodes and construct parent/child relationships. There’s a stub method for finding cousins:

class Node(object):
    """Doubly-linked node in a tree.

        >>> na = Node("na")
        >>> nb1 = Node("nb1")
        >>> nb2 = Node("nb2")

        >>> nb1.set_parent(na)
        >>> nb2.set_parent(na)

        >>> na.children
        [<Node nb1>, <Node nb2>]

        >>> nb1.parent
        <Node na>
    """

    parent = None

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

    def __repr__(self):
        return "<Node %s>" % self.data

    def set_parent(self, parent):
        """Set parent of this node.

        Also sets the children of the parent to include this node.
        """

        self.parent = parent
        parent.children.append(self)

    def cousins(self):
        """Find nodes on the same level as this node."""

Implement this method.