advent-of-code-2019/lib.py

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