Data.IntervalSet

1 files changed, 129 insertions, 0 deletions

diff --git a/src/Data/IntervalSet.hs b/src/Data/IntervalSet.hs
new file mode 100644
index 00000000..f1205274
--- /dev/null
+++ b/src/Data/IntervalSet.hs
@@ -0,0 +1,129 @@
+module Data.IntervalSet
+    ( IntSet
+    , null
+    , empty
+    , insert
+    , delete
+    , interval
+    , toIntervals
+    , nearestOutsider
+    , Data.IntervalSet.lookup
+    ) where
+
+import Prelude hiding (null)
+import qualified Data.IntMap.Strict as IntMap
+         ;import Data.IntMap.Strict (IntMap)
+import qualified Data.List as List
+import Data.Ord
+
+
+-- A set of integers.
+newtype IntSet = IntSet (IntMap Interval)
+ deriving Show
+
+-- Note: the intervalMin is not stored here but is the lookup key in an IntMap.
+data Interval = Interval
+    { intervalMax  :: {-# UNPACK #-} !Int -- ^ Maximum value contained in this interval.
+    , intervalNext :: {-# UNPACK #-} !Int -- ^ Minimum value in next interval if there is one.
+    }
+ deriving Show
+
+null :: IntSet -> Bool
+null (IntSet m) = IntMap.null m
+
+empty :: IntSet
+empty = IntSet IntMap.empty
+
+
+insert :: Int -> IntSet -> IntSet
+insert x (IntSet m) = IntSet $ case IntMap.lookupLE x m of
+    Just (lb,Interval mx ub)
+        | x <= mx   -> m
+        | otherwise -> case ub == maxBound of
+
+            True    | x == mx + 1 -> IntMap.insert lb (Interval x maxBound) m
+                    | otherwise   -> IntMap.insert lb (Interval mx x)
+                                     $ IntMap.insert x (Interval x maxBound) m
+
+            False   | mx + 2 == ub -> let (Just v', m')
+                                            = IntMap.updateLookupWithKey (\_ _ -> Nothing) ub m
+                                      in IntMap.insert lb v' m'
+                    | mx + 1 == x  -> IntMap.insert lb (Interval x ub) m
+                    | otherwise -> IntMap.insert lb (Interval mx x)
+                                   $ if ub == x + 1
+                                        then let (Just v', m')
+                                                   = IntMap.updateLookupWithKey
+                                                        (\_ _ -> Nothing) ub m
+                                             in IntMap.insert x v' m'
+                                        else IntMap.insert x (Interval x ub) m
+
+    Nothing     -> case IntMap.minViewWithKey m of
+
+        Just ((ub,v),m')
+            | x + 1 == ub -> IntMap.insert x v m'
+            | otherwise   -> IntMap.insert x (Interval x ub) m
+
+        Nothing           -> IntMap.singleton x (Interval x maxBound)
+
+member :: Int -> IntSet -> Bool
+member x (IntSet m) = case IntMap.lookupLE x m of
+    Just (lb,Interval mx _) -> x <= mx
+    Nothing                 -> False
+
+nearestOutsider :: Int -> IntSet -> Maybe Int
+nearestOutsider x (IntSet m)
+    | List.null xs = Nothing -- There are no integers outside the set!
+    | otherwise    = Just $ List.minimumBy (comparing (\y -> abs (x - y))) xs
+ where
+    xs = case IntMap.lookupLE x m of
+            Nothing -> [x]
+            Just (lb,Interval mx ub)
+             -> if ub < maxBound
+                    then case () of
+                     () | x > mx        -> [x]
+                        | minBound < lb -> [lb-1, mx+1, ub-1]
+                        | otherwise     -> [mx+1, ub-1]
+                    else case () of
+                     () | x > mx                         -> [x]
+                        | minBound < lb && mx < maxBound -> [lb-1, mx+1]
+                        | minBound < lb                  -> [lb-1]
+                        | mx < maxBound                  -> [mx+1]
+                        | otherwise                      -> []
+
+-- Note this could possibly benefit from a intervalPrev field.
+delete :: Int -> IntSet -> IntSet
+delete x (IntSet m) = IntSet $ case IntMap.lookupLE x m of
+    Nothing -> m
+    Just (lb,Interval mx nxt) -> case compare x mx of
+
+        GT -> m
+
+        EQ | lb < mx   -> IntMap.insert lb (Interval (mx - 1) nxt) m
+           | otherwise -> case IntMap.lookupLE (x-1) m of -- no intervalPrev
+                    Just (lb',Interval mx' _) -> IntMap.insert lb' (Interval mx' nxt)
+                                                 $ IntMap.delete lb m
+                    Nothing                   -> IntMap.delete lb m
+
+        LT -> let m' = IntMap.insert (x+1) (Interval mx nxt) m
+              in if lb < x
+                then IntMap.insert lb (Interval (x - 1) (x+1)) m'
+                else if x == minBound
+                    then IntMap.delete minBound m'
+                    else case IntMap.lookupLE (x-1) m' of -- no intervalPrev
+                        Just (lb',Interval mx' _) -> IntMap.insert lb' (Interval mx' (x+1))
+                                                     $ IntMap.delete lb m'
+                        Nothing                   -> IntMap.delete lb m'
+
+toIntervals :: IntSet -> [(Int,Int)]
+toIntervals (IntSet m) = List.map (\(lb,(Interval mx _)) -> (lb,mx))
+                                    $ IntMap.toList m
+
+interval :: Int -> Int -> IntSet
+interval lb mx
+    | lb <= mx  = IntSet $ IntMap.singleton lb (Interval mx maxBound)
+    | otherwise = IntSet IntMap.empty
+
+lookup :: Int -> IntSet -> Maybe (Int,Int)
+lookup k (IntSet m) = case IntMap.lookupLE k m of
+    Nothing                 -> Nothing
+    Just (lb,Interval mx _) -> Just (lb,mx)