beam: add merge_intervals utility with tests

wizard-work/beam/merge_intervals.py

@@ -0,0 +1,24 @@
+"""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`.
+"""
+from __future__ import annotations
+from typing import Iterable, Tuple, List
+
+
+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:
+        return []
+    out: List[List[int]] = [list(items[0])]
+    for s, e in items[1:]:
+        last = out[-1]
+        threshold = last[1] + (1 if merge_touching else 0)
+        if s <= threshold:
+            if e > last[1]:
+                last[1] = e
+        else:
+            out.append([s, e])
+    return out

wizard-work/beam/test_merge_intervals.py

@@ -0,0 +1,33 @@
+from merge_intervals import merge_intervals
+
+def test_empty():
+    assert merge_intervals([]) == []
+
+def test_single():
+    assert merge_intervals([(1,3)]) == [[1,3]]
+
+def test_overlap():
+    assert merge_intervals([(1,3),(2,4)]) == [[1,4]]
+
+def test_disjoint():
+    assert merge_intervals([(1,2),(4,5)]) == [[1,2],[4,5]]
+
+def test_nested():
+    assert merge_intervals([(1,10),(2,3),(4,5)]) == [[1,10]]
+
+def test_touching_default_not_merged():
+    assert merge_intervals([(1,2),(3,4)]) == [[1,2],[3,4]]
+
+def test_touching_merged():
+    assert merge_intervals([(1,2),(3,4)], merge_touching=True) == [[1,4]]
+
+def test_unsorted():
+    assert merge_intervals([(5,6),(1,2),(2,3)]) == [[1,3],[5,6]]
+
+def test_dup():
+    assert merge_intervals([(1,2),(1,2)]) == [[1,2]]
+
+if __name__=="__main__":
+    for n in [x for x in globals() if x.startswith("test_")]:
+        globals()[n]()
+    print("ALL TESTS PASSED")