The breadth-first search example.py

"""
Question:

Suppose I had a network of switches, or more generally, Nodes. Each Node
has a method getNeighbors(), which returns a set of adjacent Nodes in the
network. Write a function findNeighborhood(Node, n) which given a Node and
int n returns all nodes no more than n hops away.

Python 3.x

"""


class Node(object):
    """
    The class that represents a node of a graph
    It doesn't have all the node properties and methods
    but finding neighboring vertices.
    """

    def __init__(self, name):
        """
        A node has its name.
        We also initialize an empty list for neighboring vertices.
        """
        self.name = name
        self._neighbors = []

    def __repr__(self):
        """The simplest way to output a node to a console."""
        return "{name}({neighbors})".format(
            name=self.name,
            neighbors=','.join([node.name for node in self._neighbors]))

    def find_neighbors(self):
        """Getter for ``self._neighbors``"""
        return self._neighbors

    def got_neighbor(self, node):
        """
        The connectivy logic.
        Append a node to the neighbors of the self-node and vice versa.
        """
        self._neighbors.append(node)
        node._neighbors.append(self)

    def find_neighborhood(self, n):
        """
        Return a set of all the vertices
        that are accessable in ``n`` forward hops.

        Params:
            ``n`` - An integer (>0) number of hops to look up.

        A simple implamantation of breadth-first search
        in a graph of node connections.
        """
        neighborhood = [self]
        neighbors = self._neighbors
        for i in range(1, n):
            _neighbors = []
            for node in neighbors:
                if node in neighborhood:
                    continue
                neighborhood.append(node)
                _neighbors.extend([n for n in node._neighbors
                                  if n not in neighbors
                                  and n not in neighborhood])
            neighbors = _neighbors
        neighborhood.extend(neighbors)
        neighborhood.remove(self)
        return set(neighborhood)

    def getNeighbors(self):
        """
        Wrapper for ``self.find_neighbors``.
        The compatibility with the task.
        """
        return self.find_neighbors()

def findNeighborhood(node, n):
    """
    Wrapper for ``self.find_neighborhood``.
    The compatibility with the task.
    """
    return node.find_neighborhood(n)

if __name__ == '__main__':

    # Let's emulate a node network and run a few tests through it

    a = Node('A')
    b = Node('B')
    c = Node('C')
    d = Node('D')
    e = Node('E')
    f = Node('F')
    g = Node('G')

    a.got_neighbor(b)
    a.got_neighbor(c)
    b.got_neighbor(c)
    d.got_neighbor(c)
    e.got_neighbor(b)
    f.got_neighbor(e)
    g.got_neighbor(b)
    
    #   A
    #  / \
    # C - B - G
    # |   |
    # D   E - F

    assert findNeighborhood(a, 1) == set([c, b])
    assert findNeighborhood(a, 2) == set([c, b, g, e, d])
    assert findNeighborhood(g, 2) == set([b, a, e, c])
    assert findNeighborhood(f, 1) == set([e])
    assert findNeighborhood(g, 10) == set([a, b, c, d, e, f])

    print('All cases passed succesfully.')