Add 'between' to 'Data.Map' and 'Data.Set'

containers-tests/tests/map-properties.hs

@@ -214,6 +214,7 @@ main = defaultMain $ testGroup "map-properties"
          , testProperty "deleteMax"            prop_deleteMaxModel
          , testProperty "filter"               prop_filter
          , testProperty "partition"            prop_partition
+         , testProperty "between"              prop_between
          , testProperty "map"                  prop_map
          , testProperty "fmap"                 prop_fmap
          , testProperty "mapkeys"              prop_mapkeys
@@ -1397,6 +1398,12 @@ prop_splitAt n xs = valid taken .&&.
   where
     (taken, dropped) = splitAt n xs
 
+prop_between :: Int -> Int -> [(Int, Int)] -> Property
+prop_between lo hi xs' = valid tw .&&. tw === (filterWithKey (\k _ -> lo <= k && k <= hi) xs)
+    where
+    xs = fromList xs'
+    tw = between lo hi xs
+
 prop_takeWhileAntitone :: [(Either Int Int, Int)] -> Property
 prop_takeWhileAntitone xs' = valid tw .&&. (tw === filterWithKey (\k _ -> isLeft k) xs)
   where

containers-tests/tests/set-properties.hs

@@ -100,6 +100,7 @@ main = defaultMain $ testGroup "set-properties"
                    , testProperty "takeWhileAntitone"    prop_takeWhileAntitone
                    , testProperty "dropWhileAntitone"    prop_dropWhileAntitone
                    , testProperty "spanAntitone"         prop_spanAntitone
+                   , testProperty "between"              prop_between
                    , testProperty "take"                 prop_take
                    , testProperty "drop"                 prop_drop
                    , testProperty "splitAt"              prop_splitAt
@@ -672,6 +673,12 @@ prop_spanAntitone xs' = valid tw .&&. valid dw
     xs = fromList xs'
     (tw, dw) = spanAntitone isLeft xs
 
+prop_between :: Int -> Int -> [Int] -> Property
+prop_between lo hi xs' = valid tw .&&. tw === (filter (\x -> lo <= x && x <= hi) xs)
+    where
+    xs = fromList xs'
+    tw = between lo hi xs
+
 prop_powerSet :: Set Int -> Property
 prop_powerSet xs = valid ps .&&. ps === ps'
   where

containers/src/Data/Map/Internal.hs

@@ -299,6 +299,7 @@ module Data.Map.Internal (
     , withoutKeys
     , partition
     , partitionWithKey
+    , between
 
     , mapMaybe
     , mapMaybeWithKey
@@ -2983,6 +2984,19 @@ spanAntitone p0 m = toPair (go p0 m)
       | p kx = let u :*: v = go p r in link kx x l u :*: v
       | otherwise = let u :*: v = go p l in u :*: link kx x v r
 
+-- | /O(log n)/. Filter a map such that its keys lie between a lower and an upper bound (inclusive).
+--
+-- > between 2 4 (fromList [(3, "a"), (5, "b"), (4, "c"), (1, "d") (2, "e")]) == fromList [(3, "a"), (4, "c"), (2, "e")]
+-- > between 4 2 (fromList [(3, "a"), (5, "b"), (4, "c"), (1, "d") (2, "e")]) == empty
+-- > between 2 2 (fromList [(3, "a"), (5, "b"), (4, "c"), (1, "d") (2, "e")]) == singleton 2 "e"
+
+between :: (Ord k) => k -> k -> Map k a -> Map k a
+between _ _ Tip = Tip
+between lo hi (Bin _ kx x l r)
+    | kx < lo = between lo hi r
+    | kx > hi = between lo hi l
+    | otherwise = link kx x (dropWhileAntitone (< lo) l) (takeWhileAntitone (<= hi) r)
+
 -- | /O(n)/. Partition the map according to a predicate. The first
 -- map contains all elements that satisfy the predicate, the second all
 -- elements that fail the predicate. See also 'split'.

containers/src/Data/Map/Lazy.hs

@@ -222,6 +222,7 @@ module Data.Map.Lazy (
     , withoutKeys
     , partition
     , partitionWithKey
+    , between
     , takeWhileAntitone
     , dropWhileAntitone
     , spanAntitone

containers/src/Data/Map/Strict.hs

@@ -238,6 +238,7 @@ module Data.Map.Strict
     , withoutKeys
     , partition
     , partitionWithKey
+    , between
 
     , takeWhileAntitone
     , dropWhileAntitone

containers/src/Data/Map/Strict/Internal.hs

@@ -256,6 +256,7 @@ module Data.Map.Strict.Internal
     , takeWhileAntitone
     , dropWhileAntitone
     , spanAntitone
+    , between
 
     , mapMaybe
     , mapMaybeWithKey
@@ -393,6 +394,7 @@ import Data.Map.Internal
   , null
   , partition
   , partitionWithKey
+  , between
   , restrictKeys
   , size
   , spanAntitone

containers/src/Data/Set.hs

@@ -123,6 +123,7 @@ module Data.Set (
             , split
             , splitMember
             , splitRoot
+            , between
 
             -- * Indexed
             , lookupIndex

containers/src/Data/Set/Internal.hs

@@ -168,6 +168,7 @@ module Data.Set.Internal (
             , split
             , splitMember
             , splitRoot
+            , between
 
             -- * Indexed
             , lookupIndex
@@ -1535,6 +1536,17 @@ spanAntitone p0 m = toPair (go p0 m)
       | p x = let u :*: v = go p r in link x l u :*: v
       | otherwise = let u :*: v = go p l in u :*: link x v r
 
+-- | /O(log n)/. Filter a map set that its values lie between a lower and upper bound (inclusive).
+--
+-- > between 2 4 (fromList [3, 5, 4, 1 2]) == fromList [3, 4, 2]
+-- > between 4 2 (fromList [3, 5, 4, 1 2]) == empty
+-- > between 2 2 (fromList [3, 5, 4, 1 2]) == singleton 2
+between :: (Ord a) => a -> a -> Set a -> Set a
+between _ _ Tip = Tip
+between lo hi (Bin _ x l r)
+    | x < lo = between lo hi r
+    | x > hi = between lo hi l
+    | otherwise = link x (dropWhileAntitone (< lo) l) (takeWhileAntitone (<= hi) r)
 
 {--------------------------------------------------------------------
   Utility functions that maintain the balance properties of the tree.