95 lines
2.4 KiB
Python
95 lines
2.4 KiB
Python
#!/usr/bin/env python3
|
|
|
|
def memoize(f):
|
|
cache = dict()
|
|
def memf(*args):
|
|
key = tuple(args)
|
|
if key not in cache:
|
|
cache[key] = f(*args)
|
|
return cache[key]
|
|
return memf
|
|
|
|
|
|
class Graph(dict):
|
|
nodes = dict.keys
|
|
|
|
def add_node(self, a):
|
|
if a not in self:
|
|
self[a] = []
|
|
|
|
def add_edge(self, a, b):
|
|
self.add_node(a)
|
|
self.add_node(b)
|
|
self[a].append(b)
|
|
self[b].append(a)
|
|
|
|
def nodes_of(self, a):
|
|
return self[a]
|
|
|
|
def rm_loops(self):
|
|
for a in self.keys():
|
|
for i, b in enumerate(self.nodes_of(a)):
|
|
if a == b:
|
|
del self[a][i]
|
|
|
|
def min_path(self, a, b):
|
|
weights = dict()
|
|
visited = set()
|
|
|
|
node, w = a, 0
|
|
while node != b:
|
|
for c in self[node]:
|
|
if c in visited:
|
|
continue
|
|
if c in weights and weights[c] <= w + 1:
|
|
continue
|
|
weights[c] = w + 1
|
|
node, w = min(weights.items() , key=lambda x: x[1])
|
|
del weights[node]
|
|
visited.add(node)
|
|
return w
|
|
|
|
def min_paths(self, a):
|
|
weights = dict()
|
|
visited = set()
|
|
visited.add(a)
|
|
|
|
node, w = a, 0
|
|
while len(visited) != len(self):
|
|
for c in self[node]:
|
|
if c in visited:
|
|
continue
|
|
if c in weights and weights[c] <= w + 1:
|
|
continue
|
|
weights[c] = w + 1
|
|
node, w = min(filter(lambda x: x[0] not in visited, weights.items()) , key=lambda x: x[1])
|
|
visited.add(node)
|
|
return weights
|
|
|
|
|
|
|
|
class WeightedGraph(Graph):
|
|
def nodes_of(self, a):
|
|
return [b for b, _ in self[a]]
|
|
|
|
def add_edge(self, a, b, w):
|
|
self[a].append((b, w))
|
|
self[b].append((a, w))
|
|
|
|
def min_path(self, a, b):
|
|
weights = dict()
|
|
visited = set()
|
|
|
|
node, w = a, 0
|
|
while node != b:
|
|
for c, edge_w in self[node]:
|
|
if c in visited:
|
|
continue
|
|
if c in weights and weights[c] <= w + edge_w:
|
|
continue
|
|
weights[c] = w + edge_w
|
|
node, w = min(weights.items() , key=lambda x: x[1])
|
|
del weights[node]
|
|
visited.add(node)
|
|
return w
|