traverse.py

#!/usr/bin/env python

import unittest
import random

class TestSequenceFunctions(unittest.TestCase):

  def setUp(self):
    self.graph = { 'a': ['c'],
                  'b': ['a'],
                  'c': ['e'],
                  'd': ['a'],
                  'e': [
                      'a',
                      'b',
                      'c',
                      'd',
                      ],
                  }

  def test_traverse(self):
    traversed = traverse_graph(self.graph, 'a', 'a', 1)
    self.assertEqual( [ 'c', 'e', ], traversed, )

    traversed = traverse_graph(self.graph, 'a', 'a', 0)
    self.assertEqual( [ 'c', ], traversed, )

def traverse_graph(graph, starting_point, current, hops):
  destinations = graph[current]
  traversed = []
  if starting_point != current:
    traversed += current
  if hops < 0:
    return traversed
  for destination in destinations:
    traversed += traverse_graph(graph, starting_point, destination, hops-1)
  return traversed

if __name__ == '__main__':
  unittest.main()