Pathfinding Puzzle Solution

puzzle.py

"""

puzzle.py
Ryan Lam

Solution to this puzzle: 
https://scontent.fagc1-2.fna.fbcdn.net/v/t1.0-9/17265275_427963684219384_574285565933413029_n.jpg?oh=526bb9d07ce2930f3dd8aa867e185e00&oe=596EEC8C

"""

import sys
from functools import wraps

INF = float('inf')


class C:
    RED = 0
    GREEN = 1
    BLUE = 2
    
    
class T:
    HORSE = 0
    CAR = 1
    TROLLEY = 2
    BUS = 3
    
    
def memoized(f):
    _cache = {}
    
    @wraps(f)
    def wrapper(*args, **kwargs):
        key = (args, frozenset(kwargs.items()))
        
        if key not in _cache:
            _cache[key] = f(*args, **kwargs)
            
        assert key in _cache
        
        return _cache[key]
        
    return wrapper
    
    
class Graph:
    def __init__(self):
        self._nodes = {}
        self._edges = {}
        
    def add_node(self, id):
        self._nodes[id] = Node(id)
        
    def add_edge(self, id1, id2, color, transit):
        node1 = self.get_node(id1)
        node2 = self.get_node(id2)
        
        edge = Edge(node1, node2, color, transit)
        self._edges[frozenset((id1, id2))] = edge
        
        node1.edges.add(edge)
        node2.edges.add(edge)
        
    def get_node(self, id):
        return self._nodes[id]
        
    def get_edge(self, node1, node2):
        return self._edges[frozenset((node1.id, node2.id))]
        
    def find_path(self, startID, endID):
        startNode = self.get_node(startID)
        endNode = self.get_node(endID)
        
        startPath = (None, startNode)
        
        distances = {startPath : 0}
        predecessors = {}
        
        visited = set()
        
        lastPath = startPath
        lastEdge = None
        lastDistance = 0
        
        node = startNode
        
        while node is not endNode:
            visited.add(lastPath)
            
            if lastEdge is None:
                edges = frozenset(node.edges)
            else:
                edges = node.get_edges(lastEdge)
                
            # Update distances.
            for edge in edges:
                neighbor = edge.neighbor(node)
                
                newPath = (node, neighbor)
                
                if newPath not in distances:
                    distances[newPath] = INF
                    
                newDistance = distances[lastPath] + 1
                
                if newDistance < distances[newPath]:
                    distances[newPath] = newDistance
                    predecessors[newPath] = lastPath
                    
            # Choose next node to visit.
            minDistance = INF
            nextNode = None
            for path, distance in distances.items():
                if path in visited:
                    continue
                    
                if distance < minDistance:
                    minDistance = distance
                    _, nextNode = path
                    
                    lastPath = path
                    lastEdge = self.get_edge(*path)
                    
            if nextNode is None:
                return None     # There is no path to the end.
                
            node = nextNode
            
        # Determine the final path to the end.
        reverseFinalPath = []
        path = lastPath
        
        while path != startPath:
            reverseFinalPath.append(path[1])
            path = predecessors[path]
            
        reverseFinalPath.append(startNode)
        
        return reverseFinalPath[::-1]
        
    # # Terrible brute-force pathfinding algorithm below:
    # def find_path(self, startID, endID, _fromEdge=None, _visited=set()):
        # if startID == endID:
            # return [self.get_node(endID)]
            
        # startNode = self.get_node(startID)
        # endNode = self.get_node(endID)
        
        # assert startNode is not endNode
        
        # if _fromEdge is None:
            # assert len(_visited) == 0
            # edges = frozenset(startNode.edges)
        # else:
            # edges = startNode.get_edges(_fromEdge)
            
        # for edge in edges:
            # assert edge is not None
            
            # nextNode = edge.neighbor(startNode)
            
            # visitedPair = (startID, nextNode.id)
            
            # if visitedPair in _visited:
                # continue
                
            # _visited.add(visitedPair)
            # path = self.find_path(nextNode.id, endID, edge, _visited)
            # _visited.remove(visitedPair)
            
            # if path is not None:
                # path.append(startNode)
                # return path
                
        # else:
            # return None
            
            
class Node:
    def __init__(self, id):
        self.id = id
        self.edges = set()
        
    def __str__(self):
        return str(self.id)
        
    def __repr__(self):
        return 'Node({})'.format(repr(self.id))
        
    @memoized
    def get_edges(self, incomingEdge):
        return frozenset(
            edge for edge in self.edges
            if edge is not incomingEdge
                and (
                    edge.color == incomingEdge.color
                    or edge.transit == incomingEdge.transit
                )
        )
        
        
class Edge:
    def __init__(self, node1, node2, color, transit):
        self.node1 = node1
        self.node2 = node2
        self.color = color
        self.transit = transit
        
    def neighbor(self, node):
        if node is self.node1:
            return self.node2
        elif node is self.node2:
            return self.node1
        else:
            assert False
            
            
def make_graph():
    g = Graph()

    for id in 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghij':
        g.add_node(id)
        
    g.add_edge('A', 'B', C.RED, T.CAR)
    g.add_edge('B', 'C', C.BLUE, T.TROLLEY)
    g.add_edge('B', 'E', C.BLUE, T.CAR)
    g.add_edge('B', 'F', C.RED, T.TROLLEY)
    g.add_edge('C', 'D', C.GREEN, T.TROLLEY)
    g.add_edge('C', 'G', C.RED, T.CAR)
    g.add_edge('D', 'G', C.BLUE, T.BUS)
    g.add_edge('D', 'I', C.GREEN, T.BUS)
    g.add_edge('D', 'J', C.RED, T.TROLLEY)
    g.add_edge('E', 'F', C.GREEN, T.HORSE)
    g.add_edge('E', 'K', C.GREEN, T.BUS)
    g.add_edge('E', 'O', C.RED, T.CAR)
    g.add_edge('F', 'G', C.GREEN, T.HORSE)
    g.add_edge('F', 'K', C.RED, T.HORSE)
    g.add_edge('G', 'H', C.RED, T.CAR)
    g.add_edge('G', 'K', C.BLUE, T.BUS)
    g.add_edge('G', 'L', C.GREEN, T.TROLLEY)
    g.add_edge('H', 'I', C.RED, T.HORSE)
    g.add_edge('H', 'L', C.RED, T.BUS)
    g.add_edge('I', 'J', C.BLUE, T.HORSE)
    g.add_edge('I', 'M', C.GREEN, T.TROLLEY)
    g.add_edge('I', 'N', C.GREEN, T.HORSE)
    g.add_edge('J', 'N', C.RED, T.BUS)
    g.add_edge('K', 'O', C.BLUE, T.TROLLEY)
    g.add_edge('K', 'P', C.RED, T.CAR)
    g.add_edge('L', 'M', C.BLUE, T.BUS)
    g.add_edge('L', 'Q', C.GREEN, T.TROLLEY)
    g.add_edge('L', 'R', C.RED, T.HORSE)
    g.add_edge('M', 'N', C.BLUE, T.HORSE)
    g.add_edge('M', 'S', C.RED, T.TROLLEY)
    g.add_edge('N', 'T', C.RED, T.BUS)
    g.add_edge('O', 'U', C.RED, T.BUS)
    g.add_edge('O', 'V', C.BLUE, T.HORSE)
    g.add_edge('P', 'Q', C.GREEN, T.CAR)
    g.add_edge('P', 'V', C.GREEN, T.HORSE)
    g.add_edge('Q', 'R', C.BLUE, T.CAR)
    g.add_edge('Q', 'W', C.RED, T.TROLLEY)
    g.add_edge('R', 'S', C.BLUE, T.HORSE)
    g.add_edge('R', 'X', C.RED, T.TROLLEY)
    g.add_edge('S', 'T', C.BLUE, T.BUS)
    g.add_edge('S', 'X', C.GREEN, T.BUS)
    g.add_edge('S', 'Y', C.GREEN, T.CAR)
    g.add_edge('S', 'd', C.RED, T.HORSE)
    g.add_edge('T', 'Y', C.GREEN, T.TROLLEY)
    g.add_edge('U', 'V', C.GREEN, T.BUS)
    g.add_edge('U', 'Z', C.BLUE, T.BUS)
    g.add_edge('V', 'W', C.RED, T.CAR)
    g.add_edge('V', 'Z', C.RED, T.CAR)
    g.add_edge('V', 'a', C.GREEN, T.BUS)
    g.add_edge('V', 'b', C.BLUE, T.TROLLEY)
    g.add_edge('W', 'X', C.GREEN, T.CAR)
    g.add_edge('W', 'b', C.RED, T.TROLLEY)
    g.add_edge('X', 'c', C.RED, T.HORSE)
    g.add_edge('Y', 'e', C.GREEN, T.TROLLEY)
    g.add_edge('Z', 'a', C.RED, T.TROLLEY)
    g.add_edge('Z', 'f', C.BLUE, T.HORSE)
    g.add_edge('a', 'b', C.RED, T.CAR)
    g.add_edge('a', 'f', C.GREEN, T.HORSE)
    g.add_edge('b', 'c', C.GREEN, T.TROLLEY)
    g.add_edge('b', 'g', C.RED, T.TROLLEY)
    g.add_edge('b', 'h', C.BLUE, T.BUS)
    g.add_edge('c', 'd', C.GREEN, T.BUS)
    g.add_edge('c', 'h', C.RED, T.CAR)
    g.add_edge('d', 'e', C.GREEN, T.HORSE)
    g.add_edge('d', 'i', C.RED, T.HORSE)
    g.add_edge('e', 'i', C.RED, T.BUS)
    g.add_edge('f', 'g', C.BLUE, T.CAR)
    g.add_edge('g', 'h', C.GREEN, T.CAR)
    g.add_edge('h', 'i', C.RED, T.BUS)
    g.add_edge('i', 'j', C.BLUE, T.BUS)
    
    return g
    
    
def main():
    g = make_graph()
    
    path = g.find_path('A', 'j')
    
    if path is None:
        print("No path found!")
    else:
        print("Possible Route:")
        print('->'.join(str(node) for node in path))
        print("({} stops)".format(len(path)))
        
    return 0
    
    
if __name__ == '__main__':
    sys.exit(main())