ENH: Implement difference method for Interval and IntervalArray (pandas-dev#21998)

Co-authored-by: Pedro Frigolet <[email protected]>

Author: St0rmie

doc/source/whatsnew/v3.0.0.rst

@@ -46,6 +46,7 @@ Other enhancements
 - :meth:`DataFrame.pivot_table` and :func:`pivot_table` now allow the passing of keyword arguments to ``aggfunc`` through ``**kwargs`` (:issue:`57884`)
 - :meth:`Series.cummin` and :meth:`Series.cummax` now supports :class:`CategoricalDtype` (:issue:`52335`)
 - :meth:`Series.plot` now correctly handle the ``ylabel`` parameter for pie charts, allowing for explicit control over the y-axis label (:issue:`58239`)
+- Added :meth:`Interval.difference` and :meth:`IntervalArray.difference` to calculate the difference between interval-like objects (:issue:`21998`)
 - Added :meth:`Interval.intersection` and :meth:`IntervalArray.intersection` to calculate the intersection between interval-like objects (:issue:`21998`)
 - Added :meth:`Interval.union` and :meth:`IntervalArray.union` to calculate the union between interval-like objects (:issue:`21998`)
 - Restore support for reading Stata 104-format and enable reading 103-format dta files (:issue:`58554`)

pandas/_libs/interval.pyx

@@ -783,6 +783,138 @@ cdef class Interval(IntervalMixin):
             closed = "neither"
         return np.array([Interval(uleft, uright, closed=closed)], dtype=object)
 
+    def difference(self, other):
+        """
+        Return the difference between an interval and another.
+
+        The difference between two intervals are the points in the first
+        interval that are not shared with the second interval.
+
+        Parameters
+        ----------
+        other : Interval
+            Interval to which to calculate the difference.
+
+        Returns
+        -------
+        np.array
+            numpy array with two intervals if the second interval is
+            contained within the first. Array with one interval if
+            the difference only shortens the limits of the interval.
+            Empty array if the first interval is contained in the second
+            and thus there are no points left after difference.
+
+        Examples
+        --------
+        >>> i0 = pd.Interval(0, 3, closed='right')
+        >>> i1 = pd.Interval(2, 4, closed='right')
+        >>> i0.difference(i1)
+        array([Interval(0, 2, closed='right')], dtype=object)
+
+        >>> i2 = pd.Interval(5, 8, closed='right')
+        >>> i0.intersection(i2)
+        array([Interval(0, 3, closed='right')], dtype=object)
+
+        >>> i3 = pd.Interval(3, 5, closed='left')
+        >>> i0.difference(i3)
+        array([Interval(0, 3, closed='neither')], dtype=object)
+
+        >>> i4 = pd.Interval(-2, 7, closed='left')
+        >>> i0.difference(i4)
+        array([], dtype=object)
+
+        >>> i4.difference(i0)
+        array([Interval(-2, 0, closed='both') Interval(3, 7, closed='neither')],
+              dtype=object)
+        """
+        if not isinstance(other, Interval):
+            raise TypeError("`other` must be an Interval, "
+                            f"got {type(other).__name__}")
+
+        # if there is no overlap then the difference is the interval
+        if not self.overlaps(other):
+            return np.array([self], dtype=object)
+
+        # if the first interval is contained inside the other then there's no points
+        # left after the difference is applied
+        if self.left > other.left and self.right < other.right:
+            return np.array([], dtype=object)
+
+        # if the intervals limits match but the other interval has closed limits then
+        # there are no points left after the difference is applied
+        if (self.left == other.left and self.right == other.right and
+           other.closed_left and other.closed_right):
+            return np.array([], dtype=object)
+
+        # if the first interval contains the other then the difference is a union of
+        # two intervals
+        if self.left < other.left and self.right > other.right:
+            if self.closed_left and not other.closed_left:
+                closed1 = "both"
+            elif self.closed_left:
+                closed1 = "left"
+            elif not other.closed_left:
+                closed1 = "right"
+            else:
+                closed1 = "neither"
+
+            if self.closed_right and not other.closed_right:
+                closed2 = "both"
+            elif self.closed_right:
+                closed2 = "right"
+            elif not other.closed_right:
+                closed2 = "left"
+            else:
+                closed2 = "neither"
+
+            return np.array([Interval(self.left, other.left, closed1),
+                            Interval(other.right, self.right, closed2)],
+                            dtype=object)
+
+        # Define left limit
+        if self.left < other.left:
+            dleft = self.left
+            lclosed = self.closed_left
+        elif self.left > other.left:
+            dleft = other.right
+            lclosed = not other.closed_right
+        else:
+            dleft = other.right if other.closed_left else self.left
+            lclosed = False if other.closed_left else self.closed_left
+
+        # Define right limit
+        if self.right > other.right:
+            dright = self.right
+            rclosed = self.closed_right
+        elif self.right < other.right:
+            dright = other.left
+            rclosed = not other.closed_left
+        else:
+            dright = self.left if other.closed_right else other.right
+            rclosed = False if other.closed_right else self.closed_right
+
+        # if the interval only contains one point then it must be closed
+        # on both sides
+        if dleft == dright:
+            if (lclosed and self.closed_left) or (rclosed and self.closed_right):
+                return np.array([Interval(dleft, dright, closed="both")],
+                                dtype=object)
+            elif not (lclosed and rclosed):
+                return np.array([], dtype=object)
+
+        if dleft > dright:
+            return np.array([], dtype=object)
+
+        if lclosed and rclosed:
+            closed = "both"
+        elif lclosed:
+            closed = "left"
+        elif rclosed:
+            closed = "right"
+        else:
+            closed = "neither"
+        return np.array([Interval(dleft, dright, closed=closed)], dtype=object)
+
 
 @cython.wraparound(False)
 @cython.boundscheck(False)

pandas/core/arrays/interval.py

@@ -1510,6 +1510,71 @@ def union(self, other):
 
         return np.array([interval.union(other) for interval in self], dtype=object)
 
+    _interval_shared_docs["difference"] = textwrap.dedent(
+        """
+    Calculates difference between each Interval in the %(klass)s and a given
+    Interval.
+
+    The difference between two intervals are the points in the first
+    interval that are not shared with the second interval.
+
+    Parameters
+    ----------
+    other : Interval
+        Interval to which to calculate the difference.
+
+    Returns
+    -------
+    array
+        Array of arrays containing the differences between each interval and other.
+
+    See Also
+    --------
+    Interval.difference : Calculate difference between two Interval objects.
+
+    Examples
+    --------
+    %(examples)s
+
+    >>> intervals.difference(pd.Interval(0, 2, 'right'))
+    array([[Interval(0, 0, closed='both')],
+           [Interval(2, 5, closed='right')],
+           [Interval(2, 4, closed='right')]], dtype=object)
+
+    >>> intervals.difference(pd.Interval(2, 3, closed='left'))
+    array([array([Interval(0, 1, closed='right')], dtype=object),
+           array([Interval(1, 2, closed='neither'), Interval(3, 5, closed='both')],
+                 dtype=object)                                                     ,
+           array([Interval(3, 4, closed='right')], dtype=object)],
+          dtype=object)
+    """
+    )
+
+    @Appender(
+        _interval_shared_docs["difference"]
+        % {
+            "klass": "IntervalArray",
+            "examples": textwrap.dedent(
+                """\
+        >>> data = [(0, 1), (1, 5), (2, 4)]
+        >>> intervals = pd.arrays.IntervalArray.from_tuples(data)
+        >>> intervals
+        <IntervalArray>
+        [(0, 1], (1, 5], (2, 4]]
+        Length: 3, dtype: interval[int64, right]
+        """
+            ),
+        }
+    )
+    def difference(self, other):
+        if isinstance(other, (IntervalArray, ABCIntervalIndex)):
+            raise NotImplementedError
+        if not isinstance(other, Interval):
+            msg = f"`other` must be Interval-like, got {type(other).__name__}"
+            raise TypeError(msg)
+
+        return np.array([interval.difference(other) for interval in self], dtype=object)
+
     # ---------------------------------------------------------------------
 
     @property

pandas/tests/arrays/interval/test_overlaps.py

@@ -177,3 +177,47 @@ def test_union_invalid_type(self, other):
         msg = f"`other` must be Interval-like, got {type(other).__name__}"
         with pytest.raises(TypeError, match=msg):
             interval_container.union(other)
+
+
+class TestDifference:
+    def test_difference_interval_array(self):
+        interval = Interval(1, 8, "left")
+
+        tuples = [  # Intervals:
+            (1, 8),  # identical
+            (2, 4),  # nested
+            (0, 9),  # spanning
+            (4, 10),  # partial
+            (-5, 1),  # adjacent closed
+            (8, 10),  # adjacent open
+            (10, 15),  # disjoint
+        ]
+        interval_container = IntervalArray.from_tuples(tuples, "both")
+
+        expected = np.array(
+            [
+                np.array([Interval(8, 8, "both")], dtype=object),
+                np.array([], dtype=object),
+                np.array(
+                    [Interval(0, 1, "left"), Interval(8, 9, "both")], dtype=object
+                ),
+                np.array([Interval(8, 10, "both")], dtype=object),
+                np.array([Interval(-5, 1, "left")], dtype=object),
+                np.array([Interval(8, 10, "both")], dtype=object),
+                np.array([Interval(10, 15, "both")], dtype=object),
+            ],
+            dtype=object,
+        )
+        result = interval_container.difference(interval)
+        tm.assert_numpy_array_equal(result, expected)
+
+    @pytest.mark.parametrize(
+        "other",
+        [10, True, "foo", Timedelta("1 day"), Timestamp("2018-01-01")],
+        ids=lambda x: type(x).__name__,
+    )
+    def test_difference_invalid_type(self, other):
+        interval_container = IntervalArray.from_breaks(range(5))
+        msg = f"`other` must be Interval-like, got {type(other).__name__}"
+        with pytest.raises(TypeError, match=msg):
+            interval_container.difference(other)