#!/usr/bin/env python3 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