Restrict and delete interval for Data.Set

Author: AliceRixte

data-interval.cabal

@@ -65,6 +65,7 @@ Library
      Data.IntervalSet
      Data.IntegerInterval
      Data.Map.Interval
+     Data.Set.Interval
   Other-Modules:
      Data.Interval.Internal
      Data.IntegerInterval.Internal

src/Data/Set/Interval.hs

@@ -0,0 +1,152 @@
+{-# LANGUAGE LambdaCase #-}
+
+-----------------------------------------------------------------------------
+-- |
+-- Module      :  Data.Set.Interval
+-- Copyright   :  (c) Alice Rixte 2025
+-- License     :  BSD-style
+--
+-- Maintainer  :  masahiro.sakai@gmail.com
+-- Stability   :  provisional
+--
+-- Manipulate the keys of a "Set" (it works for both strict and lazy Sets) using
+-- "Interval"s and "IntervalSet"s
+-----------------------------------------------------------------------------
+
+module Data.Set.Interval
+  ( restrictInterval
+  , restrictIntervals
+  , deleteInterval
+  , deleteIntervals
+  ) where
+
+import Data.Set as Set
+import qualified Data.Interval.Internal as I
+import Data.Interval.Internal (Interval(..))
+import qualified Data.IntervalSet as IS
+import Data.IntervalSet (IntervalSet)
+
+
+-- | \(O(\log n)\). Restrict a 'Set' to the keys contained in a given
+-- 'Interval'.
+--
+-- >>> restrictInterval s i == filterKeys (\k -> Interval.member k i) s
+--
+-- [Usage:]
+--
+-- >>> s = Set.fromList [-2.5, 3.1, 5 , 8.5] :: Set Rational
+-- >>> restrictInterval s (3 <=..< 8.5)
+-- fromList [31 % 10,5 % 1]
+--
+-- [Performance:]
+-- This functions performs better than 'filterKeys' which is \(O(n)\).
+--
+restrictInterval :: Ord k => Set k -> Interval k -> Set k
+restrictInterval s = \case
+  Whole -> s
+  Empty -> Set.empty
+  Point k -> if Set.member k s then Set.singleton k else Set.empty
+  LessThan k -> fst $ Set.split k s
+  LessOrEqual k -> let (lt, eq, _) = Set.splitMember k s in
+    if eq then
+      Set.insert k lt
+    else
+      lt
+  GreaterThan k -> snd $ Set.split k s
+  GreaterOrEqual k -> let (_, eq, gt) = Set.splitMember k s in
+    if eq then
+      Set.insert k gt
+    else
+      gt
+  BothClosed lk uk ->
+    restrictInterval (restrictInterval s (I.GreaterOrEqual lk))
+                                         (I.LessOrEqual uk)
+  LeftOpen lk uk ->
+    restrictInterval (restrictInterval s (I.GreaterThan lk))
+                                         (I.LessOrEqual uk)
+  RightOpen lk uk ->
+    restrictInterval (restrictInterval s (I.GreaterOrEqual lk))
+                                         (I.LessThan uk)
+  BothOpen lk uk ->
+    restrictInterval (restrictInterval s (I.GreaterThan lk))
+                                         (I.LessThan uk)
+
+
+-- | \(O(n \log n)\).
+--
+-- Restrict a 'Set' to the keys contained in a given 'Interval'.
+--
+-- >>> restrictIntervals s i == filterKeys (\k -> IntervalSet.member k i) s
+--
+-- [Usage:]
+--
+-- >>> s = Set.fromList [-2.5, 3.1, 5 , 8.5] :: Set Rational
+-- >>> restrictIntervals s (IntervalSet.fromList [ -inf <..<= 3, 8 <..< inf])
+-- fromList [(-5) % 2,17 % 2]
+--
+-- [Performance:]
+-- In most cases, this function performs better than 'filterKeys', especially
+-- when the 'IntervalSet' contains few large intervals. However, in cases where
+-- the 'IntervalSet' consists of many small intervals and covers a significant
+-- portion of the Set, 'filterKeys' may outperform 'restrictIntervals' (see
+-- benchmark).
+--
+restrictIntervals :: Ord k => Set k -> IntervalSet k -> Set k
+restrictIntervals s is = IS.foldr f Set.empty is
+  where
+    f i acc = Set.union acc (restrictInterval s i)
+
+-- complexity proof:
+--
+-- This is not a tight bound, and a better analysis should be done.
+--
+-- Set.union is O(s log(n/s + 1)) where s <= n, and we know that all intevals
+-- within the IntervalSet are disjoint. So the complexity is, provided that the
+-- accumulator size is bigger than the current interval restriction size:
+--
+--   O( Σ s log(n/s + 1))
+--
+-- where i is the number of intervals in the IntervalSet, and n is the size of
+-- the original Set, and s is the size of the current interval restriction.
+--
+-- But log(n/s + 1) < log(n + 1) for all s >= 1, so we have
+--
+-- Σ s log(n/s + 1)
+-- < Σ s log(n + 1)
+-- = log(n + 1) Σ s
+-- < n log(n + 1)
+--
+-- since the sum of sizes of all interval restrictions is less than n.
+
+
+-- | \(O(n)\). Delete keys contained in a given 'Interval' from a 'Set'.
+--
+-- >>> deleteInterval i s == filterKeys (\k -> not (Interval.member k i)) s
+--
+-- [Usage:]
+--
+-- >>> s = Set.fromList [-2.5, 3.1, 5 , 8.5] :: Set Rational
+-- >>> deleteInterval (3 <=..< 8.5) s
+-- fromList [(-5) % 2,17 % 2]
+--
+-- [Performance:] In practice, this function performs a lot better than
+-- 'filterKeys' (see benchmark).
+--
+deleteInterval :: Ord k => Interval k -> Set k -> Set k
+deleteInterval i s = restrictIntervals s (IS.complement (IS.singleton i))
+
+-- | \(O(n \log n)\).
+--
+--  Delete keys contained in a given 'IntervalSet' from a 'Set'.
+--
+-- >>> deleteIntervals i s == filterKeys (\k -> not (IntervalSet.member k i)) s
+--
+-- [Usage:]
+--
+-- >>> s = Set.fromList [-2.5, 3.1, 5 , 8.5] :: Set Rational
+-- >>> deleteIntervals (IntervalSet.fromList [ -inf <..<= 3, 8 <..< inf]) s
+-- fromList [31 % 10,5 % 1]
+--
+-- See performance note of 'restrictIntervals'.
+deleteIntervals :: Ord k => IntervalSet k -> Set k -> Set k
+deleteIntervals is s = restrictIntervals s (IS.complement is)