129 lines
3.4 KiB
Haskell
129 lines
3.4 KiB
Haskell
{-# LANGUAGE NamedFieldPuns #-}
|
|
module Data.RangeMap (
|
|
HalfLine(..)
|
|
, RangeMap
|
|
, empty
|
|
, get
|
|
, line
|
|
, segment
|
|
, set
|
|
) where
|
|
|
|
type Color b = Maybe b
|
|
data Branch a b = Branch {
|
|
lower :: RangeMap a b
|
|
, bound :: a
|
|
, weight :: Int
|
|
, higher :: RangeMap a b
|
|
} deriving (Show, Read)
|
|
|
|
data RangeMap a b =
|
|
Color (Color b)
|
|
| Tree (Branch a b)
|
|
deriving (Show, Read)
|
|
|
|
empty :: RangeMap a b
|
|
empty = Color Nothing
|
|
|
|
line :: b -> RangeMap a b
|
|
line = Color . Just
|
|
|
|
get :: Ord a => RangeMap a b -> a -> Maybe b
|
|
get (Color color) _ = color
|
|
get (Tree (Branch {lower, bound, higher})) key =
|
|
get (if key < bound then lower else higher) key
|
|
|
|
data HalfLine a = Above a | Below a
|
|
|
|
origin :: HalfLine a -> a
|
|
origin (Above a) = a
|
|
origin (Below a) = a
|
|
|
|
set :: Ord a => HalfLine a -> Color b -> RangeMap a b -> RangeMap a b
|
|
set halfLine color (Tree branch@(Branch {lower, bound, higher})) =
|
|
if origin halfLine < bound
|
|
then
|
|
let newLower = set halfLine color lower in
|
|
Tree $ balance (branch {
|
|
lower = newLower
|
|
, weight = weight branch + 1
|
|
})
|
|
else
|
|
let newHigher = set halfLine color higher in
|
|
Tree $ balance (branch {
|
|
higher = newHigher
|
|
, weight = weight branch + 1
|
|
})
|
|
|
|
set (Above bound) color rangeMap = Tree (Branch {
|
|
lower = rangeMap
|
|
, bound
|
|
, weight = 0
|
|
, higher = Color color
|
|
})
|
|
|
|
set (Below bound) color rangeMap = Tree (Branch {
|
|
lower = Color color
|
|
, bound
|
|
, weight = 0
|
|
, higher = rangeMap
|
|
})
|
|
|
|
segment :: (a, a) -> b -> RangeMap a b
|
|
segment (from, to) b = Tree (Branch {
|
|
lower = empty
|
|
, bound = from
|
|
, weight = 1
|
|
, higher = Tree (Branch {
|
|
lower = line b
|
|
, bound = to
|
|
, weight = 0
|
|
, higher = empty
|
|
})
|
|
})
|
|
|
|
popMax :: Branch a b -> ((a, Color b), RangeMap a b)
|
|
popMax (Branch {lower, bound, higher = Color color}) = ((bound, color), lower)
|
|
popMax branch@(Branch {higher = Tree higherBranch, weight}) =
|
|
let (pair, subTree) = popMax higherBranch in
|
|
(pair, Tree (branch {higher = subTree, weight = weight - 1}))
|
|
|
|
popMin :: Branch a b -> ((a, Color b), RangeMap a b)
|
|
popMin (Branch {lower = Color color, bound, higher}) = ((bound, color), higher)
|
|
popMin branch@(Branch {lower = Tree lowerBranch, weight}) =
|
|
let (pair, subTree) = popMin lowerBranch in
|
|
(pair, Tree (branch {lower = subTree, weight = weight - 1}))
|
|
|
|
rotateLeft :: Ord a => Branch a b -> Branch a b
|
|
rotateLeft branch@(Branch {higher = Color _}) = branch
|
|
rotateLeft branch@(Branch {lower, bound, higher = Tree higherBranch}) =
|
|
let ((minBound, color), newHigher) = popMin higherBranch in
|
|
branch {
|
|
lower = set (Above bound) color lower
|
|
, bound = minBound
|
|
, higher = newHigher
|
|
}
|
|
|
|
rotateRight :: Ord a => Branch a b -> Branch a b
|
|
rotateRight branch@(Branch {lower = Color _}) = branch
|
|
rotateRight branch@(Branch {lower = Tree lowerBranch, bound, higher}) =
|
|
let ((maxBound, color), newLower) = popMax lowerBranch in
|
|
branch {
|
|
lower = newLower
|
|
, bound = maxBound
|
|
, higher = set (Below bound) color higher
|
|
}
|
|
|
|
getWeight :: RangeMap a b -> Int
|
|
getWeight (Tree (Branch {weight})) = weight
|
|
getWeight _ = 0
|
|
|
|
balance :: Ord a => Branch a b -> Branch a b
|
|
balance branch@(Branch {lower, bound, higher})
|
|
| lowerWeight - higherWeight > 1 = rotateRight branch
|
|
| higherWeight - lowerWeight > 1 = rotateLeft branch
|
|
| otherwise = branch
|
|
where
|
|
lowerWeight = getWeight lower
|
|
higherWeight = getWeight higher
|