Commit 9b761979 authored by Hans-Peter Deifel's avatar Hans-Peter Deifel 🐢

Implement sum bag data type

parent 2ccb7ef8
......@@ -33,6 +33,7 @@ library
, Data.Float.Utils
, Data.List.Utils
, Data.Text.Prettyprint
, Data.SumBag
, Copar.RefinementInterface
, Copar.Functors
, Copar.FunctorDescription
......
{-# LANGUAGE RoleAnnotations #-}
module Data.SumBag
( SumBag
, empty
, singleton
, size
, sum
, insert
, delete
) where
import Prelude hiding (sum, min)
import qualified Data.List.NonEmpty as NE
type SumBag a = Tree a
data Tree a = Leaf | Node (MetaData a) (Element a) (Tree a) (Tree a)
type role Tree nominal
data MetaData a = MetaData
{ nodeSize :: Int
, nodeSum :: a
}
data Element a = Element
{ value :: a
, multiplicity :: NE.NonEmpty a
}
empty :: SumBag a
empty = Leaf
singleton :: Monoid a => a -> SumBag a
singleton a =
node (Element a (NE.fromList [a])) Leaf Leaf
size :: SumBag a -> Int
size Leaf = 0
size (Node node _ _ _) = nodeSize node
sum :: Monoid a => SumBag a -> a
sum Leaf = mempty
sum (Node node _ _ _) = nodeSum node
insert :: (Ord a, Monoid a) => a -> SumBag a -> SumBag a
insert a Leaf = node (Element a (NE.fromList [a])) Leaf Leaf
insert a (Node _ e left right)
| a < value e = balance1 e (insert a left) right
| a > value e = balance1 e left (insert a right)
| otherwise = node (addOnce e) left right
delete :: (Ord a, Monoid a) => a -> SumBag a -> SumBag a
delete _ Leaf = Leaf
delete a (Node _ e left right)
| a < value e = balance1 e (delete a left) right
| a > value e = balance1 e left (delete a right)
| Just e' <- delOnce e = node e' left right
| otherwise = helper left right
where helper Leaf right = right
helper left Leaf = left
helper left right =
let (min, rest) = delmin right
in balance1 min left rest
-- Internal functions
-- | "Smart" constructor for Node. Will compute the meta data from its subtrees
node :: Monoid a => Element a -> Tree a -> Tree a -> Tree a
node a left right =
let nodeData = MetaData
{ nodeSize = size left + 1 + size right
, nodeSum = NE.head (multiplicity a) <> sum left <> sum right
}
in Node nodeData a left right
rotateSingleLeft :: Monoid a => Element a -> Tree a -> Tree a -> Tree a
rotateSingleLeft a x (Node _ b y z) = node b (node a x y) z
rotateSingleLeft _ _ _ = error "rotateSingleLeft called with empty right tree"
rotateSingleRight :: Monoid a => Element a -> Tree a -> Tree a -> Tree a
rotateSingleRight b (Node _ a x y) z = node a x (node b y z)
rotateSingleRight _ _ _ = error "rotateSingleRight called with empty left tree"
rotateDoubleLeft :: Monoid a => Element a -> Tree a -> Tree a -> Tree a
rotateDoubleLeft a x (Node _ c (Node _ b y1 y2) z) = node b (node a x y1) (node c y2 z)
rotateDoubleLeft _ _ _ = error "rotateDoubleLeft called with too small left tree"
rotateDoubleRight :: Monoid a => Element a -> Tree a -> Tree a -> Tree a
rotateDoubleRight c (Node _ a x (Node _ b y1 y2)) z = node b (node a x y1) (node c y2 z)
rotateDoubleRight _ _ _ = error "rotateDoubleRight called with too small left tree"
balance1 :: Monoid a => Element a -> Tree a -> Tree a -> Tree a
balance1 a left right
-- Subtrees have only one element
| size left + size right < 2 = node a left right
-- Right subtree is too heavy
| size right > balanceBound * size left =
let Node _ _ rleft rright = right
sizeRL = size rleft
sizeRR = size rright
in if sizeRL < sizeRR then rotateSingleLeft a left right else rotateDoubleLeft a left right
-- Left subtree is too heavy
| size left > balanceBound * size right =
let Node _ _ lleft lright = left
sizeLL = size lleft
sizeLR = size lright
in if sizeLL < sizeLR then rotateSingleRight a left right else rotateDoubleRight a left right
-- No subtree is too heave, we can just form a new tree straight away
| otherwise = node a left right
addOnce :: Semigroup a => Element a -> Element a
addOnce e = let total = NE.head (multiplicity e)
in e { multiplicity = NE.cons (total <> value e) (multiplicity e) }
delOnce :: Element a -> Maybe (Element a)
delOnce e = case snd (NE.uncons (multiplicity e)) of
Nothing -> Nothing
Just rest -> Just (e { multiplicity = rest })
delmin :: Monoid a => Tree a -> (Element a, Tree a)
delmin Leaf = error "delmin: Empty tree"
delmin (Node _ e Leaf _) = (e, Leaf)
delmin (Node _ e left right) = (\left' -> balance1 e left' right) <$> delmin left
balanceBound :: Int
balanceBound = 4
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment