TicketToWin/color_graph.py at main · jhmccoy9/TicketToWin · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
from connection import Connection
from color import Color
from destination import Destination
import numpy as np
import networkx as nx
from networkx.algorithms.approximation import traveling_salesman_problem


class ColorGraph:
    _connections = list()
    _adjacency_tensor = None  # c x d x d tensor where d is the number of destinations and c is the number of colors

    def __init__(self, connections: list[Connection]):
        self._connections = connections.copy()
        self._generateAdjacencyTensor()

    def _generateAdjacencyTensor(self) -> None:
        dest_tensor = np.full((len(Color), len(Destination), len(Destination)), np.inf)
        for connection in self._connections:
            dest_tensor[connection.color][connection.dest1][connection.dest2] = (
                connection.dist
            )
            dest_tensor[connection.color][connection.dest2][connection.dest1] = (
                connection.dist
            )
        self._adjacency_tensor = dest_tensor.copy()
        # with np.printoptions(threshold=sys.maxsize):
        #     print(dest_tensor[0], dest_tensor[0].shape)

    # Dijkstra's algorithm for simply going from one destination to another, while omitting selected colors
    # requires networkx
    def generateShortestPath(
        self, dest1: Destination, dest2: Destination, colors_to_avoid: list[Color] = []
    ):
        graph = nx.Graph()
        filtered_connections = list()
        for connection in self._connections:
            if connection.color not in colors_to_avoid:
                filtered_connections.append(
                    (connection.dest1, connection.dest2, connection.dist)
                )
        graph.add_weighted_edges_from(filtered_connections)
        shortest_path = nx.dijkstra_path(graph, source=dest1, target=dest2)
        length = nx.dijkstra_path_length(graph, source=dest1, target=dest2)
        print(
            "Path from",
            dest1.name,
            "to",
            dest2.name,
            ":",
            [station.name for station in shortest_path],
            ". Length:",
            length,
        )

    # modified traveling salesman problem for generating a plan to visit a set of locations
    # dests is a list of Destinations. colorDist is a dict that says the maximum distance
    # you are willing to travel per color between destinations
    def generatePaths(
        self,
        dests: list[Destination],
        colorDists: dict[Color, int] = {color: 6 for color in Color},
    ):
        graph = nx.Graph()
        filtered_connections = list()
        for connection in self._connections:
            try:
                if connection.dist <= colorDists[connection.color]:
                    filtered_connections.append(
                        (connection.dest1, connection.dest2, connection.dist)
                    )
            except KeyError as e:
                print(
                    "Warning: invalid color",
                    connection.color.name,
                    "accessed. This is not in your colorDists",
                    e,
                )
        graph.add_weighted_edges_from(filtered_connections)
        # peform 3 algorithms to increase your odds of getting the optimal solution
        try:
            christofides_path = traveling_salesman_problem(
                graph,
                nodes=dests,
                cycle=False,
                method=nx.algorithms.approximation.christofides,
            )
            greedy_path = traveling_salesman_problem(
                graph,
                nodes=dests,
                cycle=False,
                method=nx.algorithms.approximation.greedy_tsp,
            )

            shortest_path = (
                christofides_path
                if len(christofides_path) <= len(greedy_path)
                else greedy_path
            )
        except KeyError as e:
            print(e, "You probably filtered out a necessary path")
        except Exception as e:
            print(e)
        for path in shortest_path:
            print(path.name)