Implement sum bag data type

copar.cabal

@@ -33,6 +33,7 @@ library
                      , Data.Float.Utils
                      , Data.List.Utils
                      , Data.Text.Prettyprint
+                     , Data.SumBag
                      , Copar.RefinementInterface
                      , Copar.Functors
                      , Copar.FunctorDescription

src/Data/SumBag.hs

+{-# 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