68 lines
1.7 KiB
Python
68 lines
1.7 KiB
Python
|
#!/usr/bin/env python
|
||
|
|
|
||
|
|
import enum
|
||
|
|
import itertools
|
||
|
|
import sys
|
||
|
|
from typing import NamedTuple
|
||
|
|
|
||
|
|
|
||
|
|
class Point(NamedTuple):
|
||
|
|
x: int
|
||
|
|
y: int
|
||
|
|
|
||
|
|
|
||
|
|
class Region(enum.IntEnum):
|
||
|
|
ROCKY = 0
|
||
|
|
WET = 1
|
||
|
|
NARROW = 2
|
||
|
|
|
||
|
|
|
||
|
|
def solve(input: str) -> int:
|
||
|
|
def parse(input: list[str]) -> tuple[int, Point]:
|
||
|
|
depth = input[0].removeprefix("depth: ")
|
||
|
|
target = input[1].removeprefix("target: ")
|
||
|
|
return int(depth), Point(*(int(n) for n in target.split(",")))
|
||
|
|
|
||
|
|
def compute_erosions(depth: int, target: Point) -> dict[Point, int]:
|
||
|
|
res: dict[Point, int] = {}
|
||
|
|
for x in range(0, target.x + 1):
|
||
|
|
for y in range(0, target.y + 1):
|
||
|
|
p = Point(x, y)
|
||
|
|
if p == Point(0, 0) or p == target:
|
||
|
|
res[p] = 0
|
||
|
|
elif p.y == 0:
|
||
|
|
res[p] = p.x * 16807
|
||
|
|
elif p.x == 0:
|
||
|
|
res[p] = p.y * 48271
|
||
|
|
else:
|
||
|
|
res[p] = res[Point(p.x - 1, p.y)] * res[Point(p.x, p.y - 1)]
|
||
|
|
# Go from geologic index to erosion level
|
||
|
|
res[p] += depth
|
||
|
|
res[p] %= 20183
|
||
|
|
return res
|
||
|
|
|
||
|
|
def compute_regions(depth: int, target: Point) -> dict[Point, Region]:
|
||
|
|
return {
|
||
|
|
p: Region(erosion % 3)
|
||
|
|
for p, erosion in compute_erosions(depth, target).items()
|
||
|
|
}
|
||
|
|
|
||
|
|
depth, target = parse(input.splitlines())
|
||
|
|
regions = compute_regions(depth, target)
|
||
|
|
return sum(
|
||
|
|
regions[p]
|
||
|
|
for p in map(
|
||
|
|
Point._make,
|
||
|
|
itertools.product(range(0, target.x + 1), range(target.y + 1)),
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> None:
|
||
|
|
input = sys.stdin.read()
|
||
|
|
print(solve(input))
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|