135 lines
3.1 KiB
Python
135 lines
3.1 KiB
Python
import os
|
|
import sys
|
|
|
|
|
|
def _debug(*args, **kwargs):
|
|
kwargs["file"] = sys.stderr
|
|
print(*args, **kwargs)
|
|
|
|
|
|
Coordinate = tuple[int, int]
|
|
Schematic = dict[Coordinate, str]
|
|
|
|
|
|
def parse(instr: str) -> Schematic:
|
|
res = {}
|
|
|
|
lines = instr.splitlines()
|
|
max_x = len(lines[0])
|
|
|
|
for row_n, row in enumerate(lines):
|
|
assert len(row) == max_x
|
|
for col_n, char in enumerate(row):
|
|
res[(col_n, row_n)] = char
|
|
|
|
return res
|
|
|
|
|
|
DIRECTIONS = [(-1, -1), (0, -1), (1, -1), (-1, 0), (1, 0), (-1, 1), (0, 1), (1, 1)]
|
|
|
|
|
|
def apply_coord_delta(c: Coordinate, d: Coordinate) -> Coordinate:
|
|
a, b = c
|
|
e, f = d
|
|
return a + e, b + f
|
|
|
|
|
|
def seek_digits(
|
|
sc: Schematic, start: Coordinate, delta: Coordinate
|
|
) -> tuple[str, set[Coordinate]]:
|
|
digits = ""
|
|
coords = set()
|
|
|
|
cursor = start
|
|
while True:
|
|
cursor = apply_coord_delta(cursor, delta)
|
|
val = sc.get(cursor, ".")
|
|
if not val.isdigit():
|
|
break
|
|
coords.add(cursor)
|
|
digits += val
|
|
|
|
return digits, coords
|
|
|
|
|
|
def collect_digits_around(
|
|
sc: Schematic, start: Coordinate
|
|
) -> tuple[int, set[Coordinate]]:
|
|
backward_digits, backward_coords = seek_digits(sc, start, (-1, 0))
|
|
forward_digits, forward_coords = seek_digits(sc, start, (1, 0))
|
|
|
|
return (
|
|
int("".join(reversed(backward_digits)) + sc[start] + forward_digits),
|
|
backward_coords | forward_coords | set((start,)),
|
|
)
|
|
|
|
|
|
def one(inp: str):
|
|
schematic = parse(inp)
|
|
|
|
consumed_numbers = set()
|
|
acc = 0
|
|
|
|
for coord in schematic:
|
|
if coord in consumed_numbers:
|
|
continue
|
|
|
|
char = schematic[coord]
|
|
|
|
if not char.isdigit():
|
|
continue
|
|
|
|
is_part_number = False
|
|
for delta in DIRECTIONS:
|
|
target = schematic.get(apply_coord_delta(coord, delta), ".")
|
|
if not (target.isdigit() or target == "."):
|
|
is_part_number = True
|
|
break
|
|
|
|
if is_part_number:
|
|
n, used_coords = collect_digits_around(schematic, coord)
|
|
consumed_numbers = consumed_numbers | used_coords
|
|
acc += int(n)
|
|
|
|
return acc
|
|
|
|
|
|
def two(inp: str):
|
|
schematic = parse(inp)
|
|
acc = 0
|
|
|
|
for coord in schematic:
|
|
char = schematic[coord]
|
|
|
|
if char != "*":
|
|
continue
|
|
|
|
consumed_numbers = set()
|
|
numbers = []
|
|
|
|
for delta in DIRECTIONS:
|
|
test_coord = apply_coord_delta(coord, delta)
|
|
if test_coord in consumed_numbers:
|
|
continue
|
|
if schematic.get(test_coord, ".").isdigit():
|
|
n, c = collect_digits_around(schematic, test_coord)
|
|
consumed_numbers = consumed_numbers | c
|
|
numbers.append(n)
|
|
|
|
if len(numbers) == 2:
|
|
# is gear!
|
|
x, y = numbers
|
|
acc += x * y
|
|
|
|
return acc
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) < 2 or sys.argv[1] not in ["1", "2"]:
|
|
print("Missing day argument", file=sys.stderr)
|
|
os.exit(1)
|
|
inp = sys.stdin.read().strip()
|
|
if sys.argv[1] == "1":
|
|
print(one(inp))
|
|
else:
|
|
print(two(inp))
|