forked from lanqiao-courses/python-100
-
Notifications
You must be signed in to change notification settings - Fork 0
/
039-graph.py
68 lines (52 loc) · 1.93 KB
/
039-graph.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
from enum import Enum
class State(Enum):
unvisited = 0
visiting = 1
visited = 2
class Node:
def __init__(self, key):
self.key = key
self.visit_state = State.unvisited
self.incoming_edges = 0
self.adj_nodes = {} # Key = key, val = Node
self.adj_weights = {} # Key = key, val = weight
def __repr__(self):
return str(self.key)
def __lt__(self, other):
return self.key < other.key
def add_neighbor(self, neighbor, weight=0):
if neighbor is None or weight is None:
raise TypeError('neighbor or weight cannot be None')
neighbor.incoming_edges += 1
self.adj_weights[neighbor.key] = weight
self.adj_nodes[neighbor.key] = neighbor
def remove_neighbor(self, neighbor):
if neighbor is None:
raise TypeError('neighbor cannot be None')
if neighbor.key not in self.adj_nodes:
raise KeyError('neighbor not found')
neighbor.incoming_edges -= 1
del self.adj_weights[neighbor.key]
del self.adj_nodes[neighbor.key]
class Graph:
def __init__(self):
self.nodes = {} # Key = key, val = Node
def add_node(self, key):
if key is None:
raise TypeError('key cannot be None')
if key not in self.nodes:
self.nodes[key] = Node(key)
return self.nodes[key]
def add_edge(self, source, dest, weight=0):
if source is None or dest is None:
raise KeyError('Invalid key')
if source not in self.nodes:
self.add_node(source)
if dest not in self.nodes:
self.add_node(dest)
self.nodes[source].add_neighbor(self.nodes[dest], weight)
def add_undirected_edge(self, source, dest, weight=0):
if source is None or dest is None:
raise TypeError('key cannot be None')
self.add_edge(source, dest, weight)
self.add_edge(dest, source, weight)