diff options
-rw-r--r-- | src/Data/IntervalSet.hs | 129 |
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 @@ | |||
1 | module Data.IntervalSet | ||
2 | ( IntSet | ||
3 | , null | ||
4 | , empty | ||
5 | , insert | ||
6 | , delete | ||
7 | , interval | ||
8 | , toIntervals | ||
9 | , nearestOutsider | ||
10 | , Data.IntervalSet.lookup | ||
11 | ) where | ||
12 | |||
13 | import Prelude hiding (null) | ||
14 | import qualified Data.IntMap.Strict as IntMap | ||
15 | ;import Data.IntMap.Strict (IntMap) | ||
16 | import qualified Data.List as List | ||
17 | import Data.Ord | ||
18 | |||
19 | |||
20 | -- A set of integers. | ||
21 | newtype IntSet = IntSet (IntMap Interval) | ||
22 | deriving Show | ||
23 | |||
24 | -- Note: the intervalMin is not stored here but is the lookup key in an IntMap. | ||
25 | data Interval = Interval | ||
26 | { intervalMax :: {-# UNPACK #-} !Int -- ^ Maximum value contained in this interval. | ||
27 | , intervalNext :: {-# UNPACK #-} !Int -- ^ Minimum value in next interval if there is one. | ||
28 | } | ||
29 | deriving Show | ||
30 | |||
31 | null :: IntSet -> Bool | ||
32 | null (IntSet m) = IntMap.null m | ||
33 | |||
34 | empty :: IntSet | ||
35 | empty = IntSet IntMap.empty | ||
36 | |||
37 | |||
38 | insert :: Int -> IntSet -> IntSet | ||
39 | insert x (IntSet m) = IntSet $ case IntMap.lookupLE x m of | ||
40 | Just (lb,Interval mx ub) | ||
41 | | x <= mx -> m | ||
42 | | otherwise -> case ub == maxBound of | ||
43 | |||
44 | True | x == mx + 1 -> IntMap.insert lb (Interval x maxBound) m | ||
45 | | otherwise -> IntMap.insert lb (Interval mx x) | ||
46 | $ IntMap.insert x (Interval x maxBound) m | ||
47 | |||
48 | False | mx + 2 == ub -> let (Just v', m') | ||
49 | = IntMap.updateLookupWithKey (\_ _ -> Nothing) ub m | ||
50 | in IntMap.insert lb v' m' | ||
51 | | mx + 1 == x -> IntMap.insert lb (Interval x ub) m | ||
52 | | otherwise -> IntMap.insert lb (Interval mx x) | ||
53 | $ if ub == x + 1 | ||
54 | then let (Just v', m') | ||
55 | = IntMap.updateLookupWithKey | ||
56 | (\_ _ -> Nothing) ub m | ||
57 | in IntMap.insert x v' m' | ||
58 | else IntMap.insert x (Interval x ub) m | ||
59 | |||
60 | Nothing -> case IntMap.minViewWithKey m of | ||
61 | |||
62 | Just ((ub,v),m') | ||
63 | | x + 1 == ub -> IntMap.insert x v m' | ||
64 | | otherwise -> IntMap.insert x (Interval x ub) m | ||
65 | |||
66 | Nothing -> IntMap.singleton x (Interval x maxBound) | ||
67 | |||
68 | member :: Int -> IntSet -> Bool | ||
69 | member x (IntSet m) = case IntMap.lookupLE x m of | ||
70 | Just (lb,Interval mx _) -> x <= mx | ||
71 | Nothing -> False | ||
72 | |||
73 | nearestOutsider :: Int -> IntSet -> Maybe Int | ||
74 | nearestOutsider x (IntSet m) | ||
75 | | List.null xs = Nothing -- There are no integers outside the set! | ||
76 | | otherwise = Just $ List.minimumBy (comparing (\y -> abs (x - y))) xs | ||
77 | where | ||
78 | xs = case IntMap.lookupLE x m of | ||
79 | Nothing -> [x] | ||
80 | Just (lb,Interval mx ub) | ||
81 | -> if ub < maxBound | ||
82 | then case () of | ||
83 | () | x > mx -> [x] | ||
84 | | minBound < lb -> [lb-1, mx+1, ub-1] | ||
85 | | otherwise -> [mx+1, ub-1] | ||
86 | else case () of | ||
87 | () | x > mx -> [x] | ||
88 | | minBound < lb && mx < maxBound -> [lb-1, mx+1] | ||
89 | | minBound < lb -> [lb-1] | ||
90 | | mx < maxBound -> [mx+1] | ||
91 | | otherwise -> [] | ||
92 | |||
93 | -- Note this could possibly benefit from a intervalPrev field. | ||
94 | delete :: Int -> IntSet -> IntSet | ||
95 | delete x (IntSet m) = IntSet $ case IntMap.lookupLE x m of | ||
96 | Nothing -> m | ||
97 | Just (lb,Interval mx nxt) -> case compare x mx of | ||
98 | |||
99 | GT -> m | ||
100 | |||
101 | EQ | lb < mx -> IntMap.insert lb (Interval (mx - 1) nxt) m | ||
102 | | otherwise -> case IntMap.lookupLE (x-1) m of -- no intervalPrev | ||
103 | Just (lb',Interval mx' _) -> IntMap.insert lb' (Interval mx' nxt) | ||
104 | $ IntMap.delete lb m | ||
105 | Nothing -> IntMap.delete lb m | ||
106 | |||
107 | LT -> let m' = IntMap.insert (x+1) (Interval mx nxt) m | ||
108 | in if lb < x | ||
109 | then IntMap.insert lb (Interval (x - 1) (x+1)) m' | ||
110 | else if x == minBound | ||
111 | then IntMap.delete minBound m' | ||
112 | else case IntMap.lookupLE (x-1) m' of -- no intervalPrev | ||
113 | Just (lb',Interval mx' _) -> IntMap.insert lb' (Interval mx' (x+1)) | ||
114 | $ IntMap.delete lb m' | ||
115 | Nothing -> IntMap.delete lb m' | ||
116 | |||
117 | toIntervals :: IntSet -> [(Int,Int)] | ||
118 | toIntervals (IntSet m) = List.map (\(lb,(Interval mx _)) -> (lb,mx)) | ||
119 | $ IntMap.toList m | ||
120 | |||
121 | interval :: Int -> Int -> IntSet | ||
122 | interval lb mx | ||
123 | | lb <= mx = IntSet $ IntMap.singleton lb (Interval mx maxBound) | ||
124 | | otherwise = IntSet IntMap.empty | ||
125 | |||
126 | lookup :: Int -> IntSet -> Maybe (Int,Int) | ||
127 | lookup k (IntSet m) = case IntMap.lookupLE k m of | ||
128 | Nothing -> Nothing | ||
129 | Just (lb,Interval mx _) -> Just (lb,mx) | ||