AdventOfCode2024/d16/s.py at master · ericli-splunk/AdventOfCode2024 · 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
# Youtube: https://youtu.be/TbENqu_O1zc

import argparse, math, sys, re, functools, operator, itertools, heapq
from collections import defaultdict, Counter, deque
#sys.setrecursionlimit(100000000)
#A = list(map(int, input().split()))
#T = int(input())

def read_lines(f):
	while True:
		line = f.readline()
		if not line:
			break
		assert line[-1] == '\n'
		yield line[:-1]

def main():
	parser = argparse.ArgumentParser()
	parser.add_argument('-1', '--one', action='store_true', help='Only part 1')
	parser.add_argument('-2', '--two', action='store_true', help='Only part 2')
	parser.add_argument('input_file', nargs='?')
	args = parser.parse_args()
	if args.input_file is not None:
		f = open(args.input_file)
	else:
		f = sys.stdin
	lines = list(read_lines(f))
	if not args.two:
		print(part_1(lines))
	if not args.one:
		print(part_2(lines))

def find_map(m, X, Y, c):
	for x in range(X):
		for y in range(Y):
			if m[x][y] == c:
				return (x, y)

def part_1(lines):
	s = 0
	X = len(lines)
	Y = len(lines[0])
	sx, sy = find_map(lines, X, Y, 'S')
	ex, ey = find_map(lines, X, Y, 'E')
	visited = set()
	q = [(0, (sx, sy), (0, 1))]
	while q:
		dist, (cx, cy), (dx, dy) = heapq.heappop(q)
		if ((cx, cy), (dx, dy)) in visited:
			continue
		visited.add(((cx, cy), (dx, dy)))
		if lines[cx][cy] == 'E':
			return dist
		elif lines[cx][cy] == '#':
			pass
		elif lines[cx][cy] in ['.', 'S']:
			heapq.heappush(q, (dist + 1, (cx + dx, cy + dy), (dx, dy)))
			heapq.heappush(q, (dist + 1001, (cx + dy, cy + dx), (dy, dx)))
			heapq.heappush(q, (dist + 1001, (cx - dy, cy - dx), (-dy, -dx)))
			if lines[cx][cy] == 'S':
				heapq.heappush(q, (dist + 2001, (cx - dx, cy - dy), (-dx, -dy)))
		else:
			raise ValueError
	return s

def dfs2_bad(lines, X, Y, max_dist, dist, cx, cy, dx, dy, visited):
	if dist > max_dist:
		return False
	if lines[cx][cy] == 'E':
		assert dist == max_dist
		visited.add((cx, cy))
		return True
	elif lines[cx][cy] == '#':
		return False
	elif lines[cx][cy] in ['.', 'S']:
		args = lines, X, Y, max_dist
		c1 = dfs2(*args, dist + 1, cx + dx, cy + dy, dx, dy, visited)
		c2 = dfs2(*args, dist + 1001, cx + dy, cy + dx, dy, dx, visited)
		c3 = dfs2(*args, dist + 1001, cx - dy, cy - dx, -dy, -dx, visited)
		if lines[cx][cy] == 'S':
			c4 = dfs2(*args, dist + 2001, cx - dx, cy - dy, -dx, -dy, visited)
		else:
			c4 = False
		if any([c1, c2, c3, c4]):
			visited.add((cx, cy))
			return True
		else:
			return False
	else:
		raise ValueError

def part_2_bad(lines):
	s = 0
	X = len(lines)
	Y = len(lines[0])
	max_dist = part_1(lines)
	sx, sy = find_map(lines, X, Y, 'S')
	ex, ey = find_map(lines, X, Y, 'E')
	visited = set()
	dfs2(lines, X, Y, max_dist, 0, sx, sy, 0, 1, visited)
	s = len(visited)
	return s

def part_2(lines):
	s = 0
	X = len(lines)
	Y = len(lines[0])
	max_dist = part_1(lines)
	sx, sy = find_map(lines, X, Y, 'S')
	ex, ey = find_map(lines, X, Y, 'E')
	visited = {}
	seats = set()
	q = [(0, (sx, sy), (0, 1), [(sx, sy)])]
	while q:
		dist, (cx, cy), (dx, dy), hist = heapq.heappop(q)
		if ((cx, cy), (dx, dy)) in visited:
			if visited[((cx, cy), (dx, dy))] < dist:
				continue
			else:
				assert visited[((cx, cy), (dx, dy))] == dist
		else:
			visited[((cx, cy), (dx, dy))] = dist
		if dist > max_dist:
			continue
		hist = hist + [(cx, cy)]
		if lines[cx][cy] == 'E':
			for i in hist:
				seats.add(i)
		elif lines[cx][cy] == '#':
			pass
		elif lines[cx][cy] in ['.', 'S']:
			heapq.heappush(q, (dist + 1, (cx + dx, cy + dy), (dx, dy), hist))
			heapq.heappush(q, (dist + 1001, (cx + dy, cy + dx), (dy, dx), hist))
			heapq.heappush(q, (dist + 1001, (cx - dy, cy - dx), (-dy, -dx), hist))
			if lines[cx][cy] == 'S':
				heapq.heappush(q, (dist + 2001, (cx - dx, cy - dy), (-dx, -dy), hist))
		else:
			raise ValueError
	s = len(seats)
	return s

if __name__ == '__main__':
	main()