119 lines
3.0 KiB
Python
119 lines
3.0 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import itertools
|
|
from collections import deque
|
|
|
|
# Iterator recipes ----------------------------------------
|
|
def last(iterator):
|
|
return deque(iterator, maxlen=1)[0]
|
|
|
|
class defaultlist(list):
|
|
"""Default list class. Allows writing and reading out of bounds."""
|
|
def __init__(self, lst, val_factory):
|
|
super().__init__(lst)
|
|
self.val_factory = val_factory
|
|
|
|
def __setitem__(self, i, x):
|
|
for _ in range((i - len(self) + 1)):
|
|
self.append(self.val_factory())
|
|
super().__setitem__(i, x)
|
|
|
|
def __getitem__(self, i):
|
|
if i >= len(self):
|
|
return self.val_factory()
|
|
return super().__getitem__(i)
|
|
|
|
|
|
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
|