/
__init__.py
127 lines (98 loc) · 3.5 KB
/
__init__.py
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
from enum import Enum
from typing import *
from aocpy import BaseChallenge
class State(Enum):
EMPTY = None
WALL = 1
SAND = 2
Vector = Tuple[int, int]
Scan = Dict[Vector, State]
def parse(instr: str) -> Tuple[Scan, Vector, Vector]:
res = {}
for line in instr.strip().splitlines():
points: List[Vector] = []
for x in line.split("->"):
p = x.split(",")
points.append((int(p[0]), int(p[1])))
for i in range(len(points) - 1):
next_i = i + 1
dx = points[next_i][0] - points[i][0]
dy = points[next_i][1] - points[i][1]
# If either dx or dy is positive, that means its going down or left respectively
assert dx == 0 or dy == 0
if dx == 0:
f = lambda x, y: x + y
if dy < 0:
f = lambda x, y: x - y
for j in range(abs(dy) + 1):
res[(points[i][0], f(points[i][1], j))] = State.WALL
else:
f = lambda x, y: x + y
if dx < 0:
f = lambda x, y: x - y
for j in range(abs(dx) + 1):
res[(f(points[i][0], j), points[i][1])] = State.WALL
keys = res.keys()
min_x = min(keys, key=lambda x: x[0])[0]
max_x = max(keys, key=lambda x: x[0])[0]
min_y = min(keys, key=lambda x: x[1])[1]
max_y = max(keys, key=lambda x: x[1])[1]
return res, (min_x, min_y), (max_x, max_y)
class Challenge(BaseChallenge):
@staticmethod
def one(instr: str) -> int:
inp, min_pos, max_pos = parse(instr)
cursor = (500, 0)
grains = 0
while (min_pos[0] <= cursor[0] <= max_pos[0]) and (
0 <= cursor[1] <= max_pos[1]
):
x, y = cursor
if inp.get((x, y + 1)) is None:
cursor = (x, y + 1)
elif inp.get((x, y + 1)) == State.WALL or inp.get((x, y + 1)) == State.SAND:
if inp.get((x - 1, y + 1)) is None:
cursor = (x - 1, y + 1)
elif inp.get((x + 1, y + 1)) is None:
cursor = (x + 1, y + 1)
else:
inp[cursor] = State.SAND
grains += 1
cursor = (500, 0)
else:
inp[cursor] = State.SAND
grains += 1
cursor = (500, 0)
return grains
@staticmethod
def two(instr: str) -> int:
inp, _, max_pos = parse(instr)
max_pos = (max_pos[0], max_pos[1] + 2)
cursor = (500, 0)
grains = 0
def get(k: Vector):
if k[1] == max_pos[1]:
return State.WALL
return inp.get(k)
while True:
x, y = cursor
if get((x, y + 1)) is None:
cursor = (x, y + 1)
elif get((x, y + 1)) == State.WALL or get((x, y + 1)) == State.SAND:
if get((x - 1, y + 1)) is None:
cursor = (x - 1, y + 1)
elif get((x + 1, y + 1)) is None:
cursor = (x + 1, y + 1)
else:
inp[cursor] = State.SAND
grains += 1
if cursor == (500, 0):
break
cursor = (500, 0)
else:
inp[cursor] = State.SAND
grains += 1
if cursor == (500, 0):
break
cursor = (500, 0)
return grains