AOC 2023, day 5, part 2

nightmare.py

class Interval:
    __slots__ = ("start", "end")
    def __init__(self, start, end):
        self.start = start
        self.end = end

    def __add__(self, n):
        return Interval(self.start + n, self.end + n)

    def __repr__(self):
        return f"Interval({self.start}, {self.end})"
    
    def overlaps(self, other):
        return max(self.start, other.start) <= min(self.end, other.end)
    
    def apply_map(self, other, offset):
        if not self.overlaps(other):
            return None, [self]
    
        split_intervals = []
        overlap = Interval(max(self.start, other.start), min(self.end, other.end)) + offset
        
        if self.start < other.start:
            split_intervals.append(Interval(self.start, min(self.end, other.start - 1)))
        if self.end > other.end:
            split_intervals.append(Interval(max(self.start, other.end + 1), self.end))
        return overlap, split_intervals
        
seeds = list(map(int, inp.splitlines(1)[0].strip().replace("seeds: ", "").split()))
intervals = set(Interval(s, s+n) for s, n in zip(seeds[::2], seeds[1::2]))
inp_maps = [[list(map(int, line.split())) for line in m.splitlines()[1:]] for m in inp.split("\n\n")[1:]]
interval_maps = [
    [
        (Interval(from_start, from_start + rlen-1), to_start - from_start)
        for to_start, from_start, rlen in mapping
    ]
    for mapping in inp_maps
]
stages = [set(Interval(s, s+n-1) for s, n in zip(seeds[::2], seeds[1::2]))]
stages.extend(set() for _ in interval_maps)

for i, mapping in enumerate(interval_maps):
    while stages[i]:
        sub_stages = [set([stages[i].pop()])]
        sub_stages.extend(set() for _ in mapping)
        
        for j, (interval_map, offset) in enumerate(mapping):
            for interval in sub_stages[j]:
                overlap, split = interval.apply_map(interval_map, offset)
                sub_stages[j + 1].update(split)
                if overlap:
                    stages[i + 1].add(overlap)  # Move to next full stage
        # No map found, move to next stage:
        stages[i+1].update(sub_stages[-1])

min(interval.start for interval in stages[-1])