from collections import OrderedDict from enum import Enum # Python 2 users: Run pip install enum34 class State(Enum): unvisited = 1 visited = 2 visiting = 3 class Node: def __init__(self, id): self.id = id self.visit_state = State.unvisited self.adjacent = OrderedDict() # key = node, val = weight def __str__(self): return str(self.id) class Graph: def __init__(self): self.nodes = OrderedDict() # key = node id, val = node def add_node(self, id): node = Node(id) self.nodes[id] = node return node def add_edge(self, source, dest, weight=0): if source not in self.nodes: self.add_node(source) if dest not in self.nodes: self.add_node(dest) self.nodes[source].adjacent[self.nodes[dest]] = weight