#!/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