advent-of-code/2023/day5.py

import sys
from collections import defaultdict
from dataclasses import dataclass
from typing import Literal, TypeAlias

MAP_ORDER = [
    "seed",
    "soil",
    "fertilizer",
    "water",
    "light",
    "temperature",
    "humidity",
    "location",
]

lines = sys.stdin.read().splitlines()

seeds: list[int] = []
maps: dict[tuple[str, str], list[tuple[int, int, int]]] = {}

# parsing
index = 2
while index < len(lines):
    l0 = lines[index]
    p1, _, p2 = l0.split("-")
    p2 = p2.split()[0].strip()

    index += 1
    maps[p1, p2] = []
    while index < len(lines) and lines[index]:
        n1, n2, n3 = lines[index].split()
        maps[p1, p2].append((int(n1), int(n2), int(n3)))
        index += 1

    index += 1


def find_location(seed: int) -> int:
    value = seed
    for map1, map2 in zip(MAP_ORDER[:-1], MAP_ORDER[1:]):
        for target, start, length in maps[map1, map2]:
            if value >= start and value < start + length:
                value = target + (value - start)
                break
    return value


def find_range(
    values: tuple[int, int], map: list[tuple[int, int, int]]
) -> list[tuple[int, int]]:
    r_start, r_length = values
    ranges: list[tuple[int, int]] = []
    print(r_start, r_length)
    for target, start, length in map:
        # start is in the range
        if start <= r_start and r_start < start + length:
            if r_start + r_length < start + length:
                ranges.append(
                    (target + r_start - start, r_length)
                )
            else:
                ranges.append(
                    (target + r_start - start, length - (r_start - start))
                )
        elif start < r_start:
            if r_start + r_length < start + length:
            ranges.append(
                (target + r_start - start, target + r_start - start + r_length)
            )
        elif start >= r_start and r_start >= start

        if r_start <= start and start < start + length:
            print(start, length, target)
            if r_start + r_length < start + length:
                ranges.append(
                    (target + (start - r_start), target + (start - r_start) + length)
                )
            else:
                ranges.append((target + (start - r_start), target + length))

    return ranges


# part 1
seeds = [int(s) for s in lines[0].split(":")[1].strip().split()]
answer_1 = min(find_location(seed) for seed in seeds)
print(f"answer 1 is {answer_1}")

# part 2
parts = lines[0].split(":")[1].strip().split()
seeds_p2 = [(int(s), int(e)) for s, e in zip(parts[::2], parts[1::2])]

for seed in range(seeds_p2[0][0], seeds_p2[0][0] + seeds_p2[0][1]):
    print(seed, find_location(seed))
print("---")

seeds_p2 = [seeds_p2[0]]
for map1, map2 in zip(MAP_ORDER[:-1], MAP_ORDER[1:]):
    seeds_p2 = [s2 for s1 in seeds_p2 for s2 in find_range(s1, maps[map1, map2])]
print(seeds_p2)

answer_2 = ...
print(f"answer 2 is {answer_2}")