adventofcode2022/day16.py at main · Heraldk/adventofcode2022 · 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
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
from dataclasses import dataclass
import re
import functools


@dataclass(frozen=True)
class Valve:
    name: str
    rate: int
    neighbours: list[str]


def parse_input(lines: list[str]) -> dict[str, Valve]:
    valves = {}
    for line in lines:
        vals = re.split(
            "Valve | has flow rate=|; tunnels lead to valves |; tunnel leads to valve ",
            line.strip(),
        )
        name = vals[1]
        rate = int(vals[2])
        neighbours = [val.strip() for val in vals[3].split(",")]
        valves[name] = Valve(name=name, rate=rate, neighbours=neighbours)
    return valves


valves: dict[str, Valve] = {}


def all_pairs_shortest_path(valves: dict[str, Valve]) -> dict[str, dict[str, int]]:
    distances = {}
    for src in valves.keys():
        distances[src] = {}

    for src in valves.values():
        for dst in src.neighbours:
            distances[src.name][dst] = 1
        distances[src.name][src.name] = 0

    for k in valves:
        for i in valves:
            for j in valves:
                if k in distances[i] and j in distances[k]:
                    if (
                        j not in distances[i]
                        or distances[i][j] > distances[i][k] + distances[k][j]
                    ):
                        distances[i][j] = distances[i][k] + distances[k][j]
    return distances


def calc_flow(
    loc: str,
    rem_time: int,
    current_flow_rate: int,
    closed_valves: list[str],
    valves: dict[str, Valve],
    distances=dict[str, dict[str, int]],
) -> int:
    max_flow = 0
    # consider each closed valve: walking over to it and opening it
    for closed_valve in closed_valves:
        time_to_open = distances[loc][closed_valve] + 1
        if rem_time >= time_to_open:
            new_closed_valves = list(closed_valves)
            new_closed_valves.remove(closed_valve)
            flow = calc_flow(
                closed_valve,
                rem_time - time_to_open,
                current_flow_rate + valves[closed_valve].rate,
                new_closed_valves,
                valves,
                distances,
            )
            # how much flows out while we're walking there and opening the valve
            flow += time_to_open * current_flow_rate
            max_flow = max(flow, max_flow)

    # also just consider sitting put and not going anywhere
    flow = rem_time * current_flow_rate
    max_flow = max(flow, max_flow)

    return max_flow


def calc_flow_part_2(
    loc: str,
    loc_other: str,
    time_to_dest_other: int,
    rem_time: int,
    current_flow_rate: int,
    closed_valves: list[str],
    valves: dict[str, Valve],
    distances=dict[str, dict[str, int]],
) -> int:
    max_flow = 0
    # consider each closed valve: walking over to it and opening it
    for closed_valve in closed_valves:
        time_to_open = distances[loc][closed_valve] + 1
        if rem_time >= time_to_open:
            new_closed_valves = list(closed_valves)
            new_closed_valves.remove(closed_valve)

            # wrinkle here is need to know which of us gets to the valve first
            if time_to_dest_other >= 0 and time_to_open > time_to_dest_other:
                flow = calc_flow_part_2(
                    loc_other,
                    closed_valve,
                    time_to_open - time_to_dest_other,
                    rem_time - time_to_dest_other,
                    current_flow_rate + valves[loc_other].rate,
                    new_closed_valves,
                    valves,
                    distances,
                )
                flow += time_to_dest_other * current_flow_rate
            else:
                flow = calc_flow_part_2(
                    closed_valve,
                    loc_other,
                    time_to_dest_other - time_to_open,
                    rem_time - time_to_open,
                    current_flow_rate + valves[closed_valve].rate,
                    new_closed_valves,
                    valves,
                    distances,
                )
                flow += time_to_open * current_flow_rate
            max_flow = max(flow, max_flow)

    # also just consider sitting put and not going anywhere
    flow = 0
    if time_to_dest_other >= 0:
        flow = calc_flow_part_2(
            loc_other,
            loc,
            -1,
            rem_time - time_to_dest_other,
            current_flow_rate + valves[loc_other].rate,
            closed_valves,
            valves,
            distances,
        )
        flow += time_to_dest_other * current_flow_rate
        max_flow = max(flow, max_flow)
    else:
        flow += rem_time * current_flow_rate
        max_flow = max(flow, max_flow)

    return max_flow


with open("day16.txt", "r") as file:
    valves = parse_input(file.readlines())
    # print(calc_flow(30, 0, dict(), "AA", valves))
    distances = all_pairs_shortest_path(valves)
    flow = calc_flow(
        "AA", 30, 0, [x.name for x in valves.values() if x.rate > 0], valves, distances
    )
    print(flow)

    flow = calc_flow_part_2(
        "AA",
        "AA",
        0,
        26,
        0,
        [x.name for x in valves.values() if x.rate > 0],
        valves,
        distances,
    )
    print(flow)