mirror of
https://github.com/donnemartin/interactive-coding-challenges.git
synced 2024-03-22 13:11:13 +08:00
43 lines
1.1 KiB
Python
43 lines
1.1 KiB
Python
from enum import Enum # Python 2 users: Run pip install enum34
|
|
|
|
|
|
class State(Enum):
|
|
unvisited = 0
|
|
visiting = 1
|
|
visited = 2
|
|
|
|
|
|
class Node:
|
|
|
|
def __init__(self, id):
|
|
self.id = id
|
|
self.visit_state = State.unvisited
|
|
self.adjacent = {} # key = node, val = weight
|
|
|
|
def __str__(self):
|
|
return str(self.id)
|
|
|
|
def add_neighbor(self, neighbor, weight=0):
|
|
self.adjacent[neighbor] = weight
|
|
|
|
|
|
class Graph:
|
|
|
|
def __init__(self):
|
|
self.nodes = {} # key = node id, val = node
|
|
|
|
def add_node(self, id):
|
|
node = Node(id)
|
|
self.nodes[id] = node
|
|
return node
|
|
|
|
def add_edge(self, id_source, id_dest, weight=0):
|
|
if id_source not in self.nodes:
|
|
self.add_node(id_source)
|
|
if id_dest not in self.nodes:
|
|
self.add_node(id_dest)
|
|
self.nodes[id_source].add_neighbor(self.nodes[id_dest], weight)
|
|
|
|
def add_undirected_edge(self, source, dest, weight=0):
|
|
self.add_edge(source, dest, weight)
|
|
self.nodes[dest].add_neighbor(self.nodes[source], weight) |