Recurse Interview (DFS problem)

DFS.py

from collections import deque

class Node:
  def __init__(self,value,children):
    self.value    = value
    self.children = children
  
  def __str__(self):
    return str(self.value)

class WorkLoad:
  def __init__(self,useDFS,work=deque()):
    self.useDFS = useDFS
    self.work = work
    
  def add(self,node):
    self.work.append(node)
  
  def take(self):
    if self.useDFS:
      return self.work.pop()
    else:
      return self.work.popleft()
  
  ## added this for testing purposes
  def __str__(self):
    return str(self.work)
      
def isInGraph(graph,value,useDFS):
  """
  takes in a graph (starting node), value, 
  and a flag to use DFS (False=BFS)
  
  returns a Node, iff Node is in graph
  """
  if graph.value == value:
    return graph
  
  visited = []
  ## stk     = [graph]
  work = WorkLoad(useDFS)
  work.add(graph)
  
  while work:
    current = work.take()
    if current.value == value:
      return current
    if current in visited:
      print("ding! " + str(current.value)) ## debug purposes
      continue
    print(current.value) ## dfs verification
    visited.append(current)
    for child in current.children:
      ## why do anything twice?
      if child not in visited:
        work.add(child)
        
"""
NOTE: When switching useDFS, re-run. For some reason, running
isInGraph does not preseerve the state of the original graph.

Example::

> str(isInGraph(aNode,"e",False))
a
b
c
d
=> 'e'
> str(isInGraph(aNode,"e",True))
a
c
g
h
b
f
=> 'e'
> str(isInGraph(aNode,"e",False))
f
g
d
a
h
b
c
=> 'e'
"""

### Example on page
hNode = Node("h",[])
dNode = Node("d",[])
eNode = Node("e",[])
fNode = Node("f",[])

gNode = Node("g",[hNode])

cNode = Node("c",[gNode])
bNode = Node("b",[dNode,eNode,fNode])
aNode = Node("a",[bNode,cNode])