|
|
@ -0,0 +1,128 @@ |
|
|
|
{-# 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 |