diff --git a/wizard-work/beam/merge_intervals.py b/wizard-work/beam/merge_intervals.py index 18d014d..3b30ae9 100644 --- a/wizard-work/beam/merge_intervals.py +++ b/wizard-work/beam/merge_intervals.py @@ -1,24 +1,51 @@ """Merge overlapping intervals. -merge_intervals(intervals) -> list of [start, end] (inclusive endpoints). -Accepts any iterable of (start, end) pairs. Sorts by start then end. -Handles empty, single, touching (adjacent) intervals per `merge_touching`. +merge_intervals(intervals, *, merge_touching=False) -> list of [start, end]. + +Intervals are (start, end) pairs with inclusive integer endpoints. Input need +not be sorted. By default only genuinely overlapping intervals are merged; +adjacent intervals (end + 1 == next.start) are merged only when +merge_touching=True. Empty/None input yields []. """ + from __future__ import annotations -from typing import Iterable, Tuple, List + +from typing import Iterable, List, Sequence, Tuple, Union + +Interval = Union[Sequence[int], Tuple[int, int]] -def merge_intervals(intervals: Iterable[Tuple[int, int]], *, merge_touching: bool = False) -> List[List[int]]: - items = sorted((tuple(x) for x in intervals), key=lambda p: (p[0], p[1])) - if not items: +def merge_intervals( + intervals: Iterable[Interval] | None = None, *, merge_touching: bool = False +) -> List[List[int]]: + if intervals is None: return [] - out: List[List[int]] = [list(items[0])] - for s, e in items[1:]: - last = out[-1] + + cleaned: List[Tuple[int, int]] = [] + for iv in intervals: + if not isinstance(iv, (list, tuple)) or len(iv) != 2: + raise ValueError(f"interval must be a 2-element sequence, got {iv!r}") + a, b = iv + if isinstance(a, bool) or isinstance(b, bool): + raise ValueError("bool is not a valid interval endpoint") + if not isinstance(a, int) or not isinstance(b, int): + raise ValueError(f"endpoints must be ints, got {a!r}, {b!r}") + if a > b: + raise ValueError(f"start must be <= end, got {a}, {b}") + cleaned.append((a, b)) + + if not cleaned: + return [] + + cleaned.sort(key=lambda p: (p[0], p[1])) + + merged: List[List[int]] = [list(cleaned[0])] + for start, end in cleaned[1:]: + last = merged[-1] threshold = last[1] + (1 if merge_touching else 0) - if s <= threshold: - if e > last[1]: - last[1] = e + if start <= threshold: + if end > last[1]: + last[1] = end else: - out.append([s, e]) - return out \ No newline at end of file + merged.append([start, end]) + return merged