advent-of-code-2019/lib.py

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