{-
 - Rød-sorte træer implementeret i Haskell
 -
 - (c) svalle@imada.sdu.dk
 -}

module RedBlackTree where

data Colour = R | B
            deriving (Show, Eq)
data Tree a = L| N Colour (Tree a) a (Tree a)
            deriving (Show, Eq)

---------------------------------------------------------------------------
-- Et tomt træ
emptyTree :: Tree a
emptyTree = L

---------------------------------------------------------------------------
-- Er et element i træet?
search :: Ord a => Tree a -> a -> Bool
search L _ = False
search (N _ l n r) v
  | v < n     = search l v
  | v > n     = search r v
  | otherwise = True

---------------------------------------------------------------------------
-- Indsættelser
-- (særligt tilfælde til roden, så den altid er sort)
ins :: Ord a => a -> Tree a -> Tree a
ins nv t = N B l v r
  where
    N _ l v r = ins' nv t

-- (generelt tilfælde)
ins' :: Ord a => a -> Tree a -> Tree a
ins' nv t@(N c l v r)
  | nv < v    = reb (N c (ins' nv l) v r)
  | nv > v    = reb (N c l v (ins' nv r))
  | otherwise = t
int' nv L = N R L nv L 

reb :: Tree a -> Tree a
-- tilfælde A.1:
reb (N B (N R (N R a x b) y c) z (N R d u e)) =
    (N R (N B (N R a x b) y c) z (N B d u e))
-- tilfælde A.2: 
reb (N B (N R a x (N R b y c)) z (N R d u e)) =
    (N R (N B a x (N R b y c)) z (N B d u e))
-- tilfælde B: Ordnes i ins
-- tilfælde C:
reb (N B (N R (N R a x b) y c) z (N B d u e)) =
    (N B (N R a x b) y (N R c z (N B d u e)))
-- tilfælde D:
reb (N B (N R a x (N R b y c)) z (N B d u e)) =
    (N B (N R a x b) y (N R c z (N B d u e)))

-- SPEJLVENDT:
-- tilfælde A.1:
reb (N B (N R a x b) y (N R c z (N R d u e))) =
    (N R (N B a x b) y (N B c z (N R d u e)))
-- tilfælde A.2: 
reb (N B (N R a x b) y (N R (N R c z d) u e)) =
    (N R (N B a x b) y (N B (N R c z d) u e))
-- tilfælde B: Ordnes i ins
-- tilfælde C:
reb (N B (N B a x b) y (N R c z (N R d u e))) =
    (N B (N R (N B a x b) y c) z (N R d u e))
-- tilfælde D:
reb (N B (N B a x b) y (N R (N R c z d) u e)) = 
    (N B (N R (N B a x b) y c) z (N R d u e))

-- Hvis intet af det andet pattern-matcher,
-- skal der ikke rebalanceres:
reb t = t

---------------------------------------------------------------------------
-- Træ-til-list-funktion
tree2list :: Tree a -> [a]
tree2list tree = t2lX tree []
  where
    t2lX L acc = acc
    t2lX (N _ t1 n t2) acc = (t2lX t1 (n : t2lX t2 acc))

---------------------------------------------------------------------------
-- Liste-til-træ-funktioner

-- For både u- og sorterede lister
list2tree :: Ord a => [a] -> Tree a
list2tree xs
  | isSorted xs = slist2tree xs
  | otherwise   = ulist2tree xs                    

isSorted :: Ord a => [a] -> Bool
isSorted [] = True
isSorted [x] = True
isSorted (x:y:xs) = x < y && isSorted (y:xs)

-- For usorterede lister
ulist2tree :: Ord a => [a] -> Tree a
ulist2tree xs = foldr ins emptyTree xs

-- For sorterede lister
slist2tree :: [a] -> Tree a
slist2tree [] = L
slist2tree xs = fst (aux xs ht lg)
  where
    lg = length xs
    ht = noBLayers lg

-- find ud af antallet af "sorte" lag i træet
-- et evt. ikke-udfyldt nederste lag vil blive "rødt"
noBLayers :: Integral a => a -> a
noBLayers n = log2 (n+1)

log2 :: Integral a => a -> a
log2 n = floor (logBase 2 (fromIntegral n))

aux :: [a] -> Int -> Int -> (Tree a,[a])
-- aux xs height use
-- base-tilfælde
aux xs         0 0 = (L,xs)
aux (x:xs)     0 1 = (N R L x L, xs)
-- generelt tilfælde
aux xs h use       = (N B lt x rt, xs3)
  where
    hlf1 = use - 1 - (div (use-1) 2)
    hlf2 = use - 1 - hlf1
    (lt,x:xs2) = aux xs  (h-1) hlf1
    (rt,xs3) = aux xs2 (h-1) hlf2
               
-- Check-funktion, der skriver den sorte højde ud for alle blade
hg t = hglist t 0
hglist L n = [n]
hglist (N B l x r) n = hglist l (n+1) ++ hglist r (n+1)
hglist (N R l x r) n = hglist l n ++ hglist r n