diff --git a/README.md b/README.md index fb2c6fb..59314ae 100644 --- a/README.md +++ b/README.md @@ -45,6 +45,8 @@ Features * `tree.remove_overlap(point)` * `tree.remove_overlap(begin, end)` (removes all overlapping the range) * `tree.remove_envelop(begin, end)` (removes all enveloped in the range) + * `tree.remove_data(data)` (removes all intervals containing this data) + * `tree.remove_data(data, point)` (removes all intervals containing this data and overlapping the point) * Point queries * `tree[point]` diff --git a/intervaltree/intervaltree.py b/intervaltree/intervaltree.py index 2548eed..e999f3d 100644 --- a/intervaltree/intervaltree.py +++ b/intervaltree/intervaltree.py @@ -149,6 +149,15 @@ class IntervalTree(MutableSet): >>> tree IntervalTree() + Delete intervals, containing certain data: + >>> tree = IntervalTree([Interval(-10, 0, "x"), Interval(0, 10, "y"), Interval(-10, 0, "y")]) + >>> tree + IntervalTree([Interval(-10, 0, "x"), Interval(0, 10, "y"), Interval(-10, 0, "y")]) + >>> tree.remove_data("y", set(1)) + IntervalTree([Interval(-10, 0, "x"), Interval(-10, 0, "y")]) + >>> tree.remove_data("y") + IntervalTree([Interval(-10, 0, "x")]) + Point queries:: >>> tree = IntervalTree([Interval(-1.1, 1.1), Interval(-0.5, 1.5), Interval(0.5, 1.7)]) @@ -467,6 +476,28 @@ def symmetric_difference_update(self, other): other.remove(iv) self.update(other) + def remove_data(self, data, points=None): + """ + Removes all intervals containing the given data. If a set of points is provided, + only intervals containing the points AND containing the given data will be removed. + Providing a set of points speeds up the operation considerably. + """ + if data is None: + raise ValueError("IntervalTree: No data submitted.") + + if points is not None: + root = self.top_node + ivs = set() + for point in points: + ivs |= root.search_point(point, set()) + ivs_with_data = [iv for iv in ivs if iv.data == data] + for iv in ivs_with_data: + self.discard(iv) + else: + ivs_without_data = [iv for iv in self.items() if iv.data != data] + self.clear() + self.__init__(ivs_without_data) + def remove_overlap(self, begin, end=None): """ Removes all intervals overlapping the given point or range.