2-3tree.hs

import Data.Maybe (fromMaybe)

data TwoThreeSet a = Empty | One a | Tree (TwoThreeTree a)

emptyMap :: TwoThreeSet a
emptyMap = Empty

search :: Ord a => TwoThreeSet a -> a -> Bool
search Empty _ = False
search (One x) y | x == y    = True
                 | otherwise = False
search (Tree t) x = tSearch t x

insert :: Ord a => TwoThreeSet a -> a -> TwoThreeSet a
insert Empty x = One x
insert one@(One x) y | x == y = one
                     | x > y  = Tree $ TwoNode (Leaf y) (Leaf x) x
                     | x < y  = Tree $ TwoNode (Leaf x) (Leaf y) y
insert (Tree t) x = Tree $ tInsert t x

delete :: Ord a => TwoThreeSet a -> a -> TwoThreeSet a
delete Empty x = Empty
delete one@(One x) y | x == y    = Empty
                     | otherwise = one
delete (Tree t) x = case tDelete t x of
                      Leaf y -> One y
                      node -> Tree node


data TwoThreeTree a  = Leaf a
                     | TwoNode (TwoThreeTree a) (TwoThreeTree a) a
                     | ThreeNode (TwoThreeTree a) (TwoThreeTree a) (TwoThreeTree a) a a
                     deriving Show

tSearch :: Ord a => TwoThreeTree a -> a -> Bool
tSearch (Leaf y) x | x == y    = True
                   | otherwise = False
tSearch (TwoNode lTree rTree rMin) x
    | x < rMin  = tSearch lTree x
    | otherwise = tSearch rTree x
tSearch (ThreeNode lTree cTree rTree cMin rMin) x
    | x < cMin  = tSearch lTree x
    | x < rMin  = tSearch cTree x
    | otherwise = tSearch rTree x

tInsert :: Ord a => TwoThreeTree a -> a -> TwoThreeTree a
tInsert t x = case tInsert_ t x of
                Only t1 -> t1
                Brother t1 t2 min -> TwoNode t1 t2 min

-- second parameter of Brother is so-called new tree (new tree is righter one of divided trees)
-- third paramater of Brother is minimal value of new tree
data InsertedSubTree a = Only (TwoThreeTree a) | Brother (TwoThreeTree a) (TwoThreeTree a) a

tInsert_ :: Ord a => TwoThreeTree a -> a -> InsertedSubTree a
tInsert_ leaf@(Leaf y) x | x == y = Only leaf
                         | x > y  = Brother leaf (Leaf x) x
                         | x < y  = Brother (Leaf x) leaf y -- inserting minimum

tInsert_ (TwoNode lTree rTree rMin) x
    | x < rMin  = refreshLeft $ tInsert_ lTree x
    | otherwise = refreshRight $ tInsert_ rTree x
  where
    refreshLeft (Only t) = Only $ TwoNode t rTree rMin
    refreshLeft (Brother t1 t2 min) = Only $ ThreeNode t1 t2 rTree min rMin
    refreshRight (Only t) = Only $ TwoNode lTree t rMin
    refreshRight (Brother t1 t2 min) = Only $ ThreeNode lTree t1 t2 rMin min

tInsert_ (ThreeNode lTree cTree rTree cMin rMin) x
    | x < cMin  = refreshLeft $ tInsert_ lTree x
    | x < rMin  = refreshCenter $ tInsert_ cTree x
    | otherwise = refreshRight $ tInsert_ rTree x
  where
    refreshLeft (Only t) = Only $ ThreeNode t cTree rTree cMin rMin
    refreshLeft (Brother t1 t2 min) = Brother (TwoNode t1 t2 min) (TwoNode cTree rTree rMin) cMin
    refreshCenter (Only t) = Only $ ThreeNode lTree t rTree cMin rMin
    refreshCenter (Brother t1 t2 min) = Brother (TwoNode lTree t1 cMin) (TwoNode t2 rTree rMin) min
    refreshRight (Only t) = Only $ ThreeNode lTree cTree t cMin rMin
    refreshRight (Brother t1 t2 min) = Brother (TwoNode lTree cTree cMin) (TwoNode t1 t2 min) rMin


-- first paramater of Keep and Orphan is minimal value of subtree
data DeletedSubTree a = Keep (Maybe a) (TwoThreeTree a) | Orphan (Maybe a) (TwoThreeTree a) | DeleteLeaf Bool

tDelete :: Ord a => TwoThreeTree a -> a -> TwoThreeTree a
tDelete t x = case tDelete_ t x of
                Keep _ node -> node
                Orphan _ node -> node
                DeleteLeaf _ -> error "deleting leaf"

tDelete_ :: Ord a => TwoThreeTree a -> a -> DeletedSubTree a
tDelete_ (Leaf y) x | x == y    = DeleteLeaf True
                    | otherwise = DeleteLeaf False

tDelete_ node@(TwoNode lTree rTree rMin) x
    | x < rMin  = refreshLeft $ tDelete_ lTree x
    | otherwise = refreshRight $ tDelete_ rTree x
  where
    refreshLeft (DeleteLeaf True) = Orphan (Just rMin) rTree
    refreshLeft (DeleteLeaf False) = Keep Nothing node
    refreshLeft (Keep maybeMin t) = Keep maybeMin $ TwoNode t rTree rMin
    refreshLeft (Orphan maybeMin t) =
      case rTree of
        (ThreeNode _ _ _ _ _) -> let (lTree', rTree', rMin') = shiftLeft t rTree rMin
                                 in  Keep maybeMin $ TwoNode lTree' rTree' rMin'
        (TwoNode _ _ _) -> Orphan maybeMin $ mergeRight t rTree rMin
    refreshRight (DeleteLeaf True) = Orphan Nothing lTree
    refreshRight (DeleteLeaf False) = Keep Nothing node
    refreshRight (Keep maybeMin t) = Keep Nothing $ TwoNode lTree rTree rMinUpdate
      where rMinUpdate = fromMaybe rMin maybeMin
    refreshRight (Orphan maybeMin t) =
      case rTree of
        (ThreeNode _ _ _ _ _) -> let (lTree', rTree', rMin') = shiftRight lTree t rMinUpdate
                                 in  Keep Nothing $ TwoNode lTree' rTree' rMin'
        (TwoNode _ _ _) -> Orphan Nothing $ mergeLeft lTree t rMinUpdate
      where rMinUpdate = fromMaybe rMin maybeMin

tDelete_ node@(ThreeNode lTree cTree rTree cMin rMin) x
    | x < cMin  = refreshLeft   $ tDelete_ lTree x
    | x < rMin  = refreshCenter $ tDelete_ cTree x
    | otherwise = refreshRight  $ tDelete_ rTree x
  where
    refreshLeft (DeleteLeaf True) = Keep (Just cMin) $ TwoNode cTree rTree rMin
    refreshLeft (DeleteLeaf False) = Keep Nothing node
    refreshLeft (Keep maybeMin t) = Keep maybeMin $ ThreeNode t cTree rTree cMin rMin
    refreshLeft (Orphan maybeMin t) =
      case rTree of
        (ThreeNode _ _ _ _ _) -> let (lTree', cTree', cMin') = shiftLeft t cTree cMin
                                 in  Keep maybeMin $ ThreeNode lTree' cTree' rTree cMin' rMin
        (TwoNode _ _ _) -> Keep maybeMin $ TwoNode  (mergeRight t cTree cMin) rTree rMin
    refreshCenter (DeleteLeaf True) = Keep Nothing $ TwoNode lTree rTree rMin
    refreshCenter (DeleteLeaf False) = Keep Nothing node
    refreshCenter (Keep maybeMin t) = Keep Nothing $ ThreeNode lTree t rTree cMinUpdate rMin
      where cMinUpdate = fromMaybe cMin maybeMin
    refreshCenter (Orphan maybeMin t) =
      case rTree of
        (ThreeNode _ _ _ _ _) -> let (cTree', rTree', rMin') = shiftLeft t rTree rMin
                                 in  Keep Nothing $ ThreeNode lTree cTree' rTree' cMinUpdate rMin'
        (TwoNode _ _ _) -> Keep Nothing $ TwoNode lTree (mergeRight t rTree rMin) cMinUpdate
      where cMinUpdate = fromMaybe cMin maybeMin
    refreshRight (DeleteLeaf True) = Keep Nothing $ TwoNode lTree cTree cMin
    refreshRight (DeleteLeaf False) = Keep Nothing node
    refreshRight (Keep maybeMin t) = Keep Nothing $ ThreeNode lTree cTree t cMin rMinUpdate
      where rMinUpdate = fromMaybe rMin maybeMin
    refreshRight (Orphan maybeMin t) =
      case rTree of
        (ThreeNode _ _ _ _ _) -> let (cTree', rTree', rMin') = shiftRight cTree t rMinUpdate
                                 in  Keep Nothing $ ThreeNode lTree cTree' rTree' cMin rMin'
        (TwoNode _ _ _) -> Keep Nothing $ TwoNode lTree (mergeLeft cTree t rMinUpdate) cMin
      where rMinUpdate = fromMaybe rMin maybeMin

shiftRight (ThreeNode lTree cTree rTree cMin rMin) t min =
    (TwoNode lTree cTree cMin, TwoNode rTree t min, rMin)
shiftLeft  t (ThreeNode lTree cTree rTree cMin rMin) min =
    (TwoNode t lTree min, TwoNode cTree rTree rMin, cMin)
mergeLeft (TwoNode lTree rTree rMin) t min =
    (ThreeNode lTree rTree t rMin min)
mergeRight t (TwoNode lTree rTree rMin) min =
    (ThreeNode t lTree rTree min rMin)