23Tree implementation in Haskell (http://en.wikipedia.org/wiki/2-3_tree)

23Tree.hs

{-# LANGUAGE InstanceSigs #-}
module Data.TwoThreeTree where

import Control.Applicative
import Data.Foldable
import Data.Traversable
import Prelude hiding (foldr)


data Tree23 k = Leaf
              | OneElem (Tree23 k) k (Tree23 k)
              | TwoElems (Tree23 k) k (Tree23 k) k (Tree23 k)
  deriving Eq

instance Show k => Show (Tree23 k) where
  show Leaf = "_"
  show (OneElem l k r) = '(':show l ++ ' ':show k ++ ' ':show r ++ ")"
  show (TwoElems l k c k' r) = '(':show l ++ ' ':show k ++ ' ':show c ++
                               ' ':show k' ++ ' ':show r ++ ")"

-- | Insert an element in the tree
--
-- Examples:
--
-- >>> insert 1 Leaf
-- (_ 1 _)
-- >>> insert 2 . insert 1 $ Leaf
-- (_ 1 _ 2 _)
-- >>> insert 0 . insert 2 . insert 1 $ Leaf
-- ((_ 0 _) 1 (_ 2 _))
-- >>> insert 10 . insert 0 . insert 2 . insert 1 $ Leaf
-- ((_ 0 _) 1 (_ 2 _ 10 _))
-- >>> let t = insert 5 . insert 10 . insert 0 . insert 2 . insert 1 $ Leaf
-- >>> t
-- ((_ 0 _) 1 (_ 2 _) 5 (_ 10 _))
-- >>> let t1 = insert 3 t
-- >>> t1
-- ((_ 0 _) 1 (_ 2 _ 3 _) 5 (_ 10 _))
-- >>> let t2 = insert 4 t1
-- >>> t2
-- (((_ 0 _) 1 (_ 2 _)) 3 ((_ 4 _) 5 (_ 10 _)))
-- >>> let t3 = insert (-1) t2
-- >>> t3
-- (((_ -1 _ 0 _) 1 (_ 2 _)) 3 ((_ 4 _) 5 (_ 10 _)))
insert :: (Ord k)
       => k
       -> Tree23 k
       -> Tree23 k
insert key tree = let (n,mn) = insert' key tree
                  in case mn of
                      Nothing -> n
                      Just (k,r) -> OneElem n k r
  where 
    insert' :: (Ord k)
            => k
            -> Tree23 k
            -- returns the new left and maybe (in case of split) the new right
            -> (Tree23 k,Maybe (k,Tree23 k))
    insert' k Leaf = (OneElem Leaf k Leaf,Nothing)
    insert' k on@(OneElem l k' r)
      | k < k' =
        if l == Leaf then (TwoElems l k  Leaf k' r,Nothing)
                    else let (nl,mnr) = insert' k l
                         in let res = case mnr of
                                  Nothing -> OneElem nl k' r
                                  Just (k'',t) -> TwoElems nl k'' t k'   r
                         in (res,Nothing)
      | k > k' =
        if r == Leaf then (TwoElems l k' Leaf k  r,Nothing)
                    else let (nr,mnl) = insert' k r
                         in let res = case mnl of
                                  Nothing -> OneElem l k' nr
                                  Just (k'',t) -> TwoElems l k' nr k'' t
                         in (res,Nothing)
     -- Element found, nothing to add
      | otherwise = (on,Nothing)
    insert' k tn@(TwoElems l lk c rk r)
      | k < lk =
          if l == Leaf
            then (OneElem l k c,Just (lk,OneElem Leaf rk r))
            else let (nl,mnr) = insert' k l
                 in case mnr of
                      Nothing -> (TwoElems nl lk c rk r,Nothing)
                      Just (k',t) -> (OneElem nl k' t
                                     ,Just (lk,OneElem c rk r))
      | k > lk && k < rk =
          if c == Leaf
            then (OneElem l lk c,Just(k,OneElem Leaf rk r))
            else let (nc,mn) = insert' k c
                 in case mn of
                      Nothing -> (TwoElems l lk nc rk r,Nothing)
                      Just (k',t) -> (OneElem l lk nc
                                    ,Just (k',OneElem t rk r))
      | k > rk =
          if r == Leaf
            then (OneElem l lk c,Just (rk,OneElem Leaf k r))
            else let (nr,mn) = insert' k r
                 in case mn of
                      Nothing -> (TwoElems l lk c rk nr,Nothing)
                      Just (k',t) -> (OneElem l lk c
                                    ,Just (rk,OneElem nr k' t))
     -- Element found, nothing to add
     | otherwise = (tn,Nothing)

-- | If the value is an element of the tree
--
-- Examples:
--
-- >>> 1 `elemOf` Leaf
-- False
-- >>> 1 `elemOf` (insert 1 Leaf)
-- True
-- >>> 1 `elemOf` (insert 2 Leaf)
-- False
-- >>> 1 `elemOf` (insert 1 . insert 2 $ Leaf)
-- True
-- >>> 1 `elemOf` (insert 2 . insert 3 $ Leaf)
-- False
-- >>> 1 `elemOf`  (insert 1 . insert 10 . insert 3 . insert 5 $ Leaf)
-- True
-- >>> 1 `elemOf`  (insert 6 . insert 10 . insert 3 . insert 5 $ Leaf)
-- False
elemOf :: (Ord k)
       => k
       -> Tree23 k
       -> Bool
k `elemOf` Leaf = False
k `elemOf` (OneElem l k' r) = case k `compare` k' of
                                LT -> k `elemOf` l
                                EQ -> True
                                GT -> k `elemOf` r
k `elemOf` (TwoElems l kl c kr r) = case k `compare` kl of
                                     LT -> k `elemOf` l
                                     EQ -> True
                                     GT -> case k `compare` kr of
                                           LT -> k `elemOf` c
                                           EQ -> True
                                           GT -> k `elemOf` r

-- Instances

-- | 23tree is a Functor
--
-- Examples:
--
-- >>> fmap (+1) $ insert 10 . insert 1 $ Leaf
-- (_ 2 _ 11 _)
instance Functor Tree23 where
  fmap :: (a -> b) -> Tree23 a -> Tree23 b
  fmap f Leaf = Leaf
  fmap f (OneElem l k r) = OneElem (fmap f l) (f k) (fmap f r)
  fmap f (TwoElems l kl c kr r) = TwoElems (fmap f l) (f kl) (fmap f c)
                                                      (f kr) (fmap f r)

-- | 23tree is Foldable
--
-- Examples:
--
-- >>> toList $ insert 10 . insert 1 . insert 5 $ Leaf
-- [1,5,10]
instance Foldable Tree23 where
  -- foldr :: (a -> b -> b) -> b -> t a -> b
  foldr :: (a -> b -> b) -> b -> Tree23 a -> b
  foldr f a Leaf = a
  foldr f a (OneElem l k r) = foldr f (f k (foldr f a r)) l
  foldr f a (TwoElems l kl c kr r) = foldr f
                                      (f kl
                                        (foldr f (f kr (foldr f a r)) c)) l

-- | 23tree is Traversable
--
-- Examples:
--
-- >>> traverse_ print $ insert 10 . insert 1 $ Leaf
-- 1
-- 10
instance Traversable Tree23 where
  -- traverse :: Applicative f => (a -> f b) -> t a -> f (t b)Source
  traverse :: Applicative f => (a -> f b) -> Tree23 a -> f (Tree23 b)
  traverse f Leaf = pure Leaf
  traverse f (OneElem l k r) = OneElem <$> traverse f l <*> f k <*> traverse f r
  traverse f (TwoElems l kl c kr r) = TwoElems <$> traverse f l
                                               <*> f kl
                                               <*> traverse f c
                                               <*> f kr
                                               <*> traverse f r