(**************************************************************************)
(*                                                                        *)
(*  Copyright (C) Richard S Bird                                          *)
(*                                                                        *)
(*  This software is free software; you can redistribute it and/or        *)
(*  modify it under the terms of the GNU Library General Public           *)
(*  License version 2.1, with the special exception on linking            *)
(*  described in file LICENSE.                                            *)
(*                                                                        *)
(*  This software is distributed in the hope that it will be useful,      *)
(*  but WITHOUT ANY WARRANTY; without even the implied warranty of        *)
(*  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.                  *)
(*                                                                        *)
(**************************************************************************)

The Garsia-Wachs algorithm:

gwa  -- the quadratic version
gwaL -- the linearithmic version

--------------------------------------------------------------------------------
1. Basics

> import Data.Array (array, elems)
> import System.Random hiding (split)

> rands n b x = take n (randomRs (1,b) (mkStdGen x)) :: [Int]

> data Tree a  = Leaf a | Fork (Tree a) (Tree a) deriving Eq
> type Weight  = Int
> type Label   = Int
> type Depth   = Int

> showTree :: Show a => Tree a -> String
> showTree (Leaf x)     = show x
> showTree (Fork e1 e2) = "(" ++ showTree e1 ++ 
>                         " " ++ showTree e2 ++ ")"

> instance Show a => Show (Tree a)
>  where show = showTree

> size :: Tree a -> Int
> size (Leaf _)   = 1
> size (Fork u v) = size u + size v

> fringe :: Tree a -> [a]
> fringe t = help t []
>  where help (Leaf x) xs   = x:xs
>        help (Fork u v) xs = help u (help v xs)

> depths :: Tree a -> [Depth]
> depths t = from 0 t []
>  where
>  from d (Leaf x) ds   = d:ds
>  from d (Fork u v) ds = from (d+1) u (from (d+1) v ds) 

> cost :: Tree Weight -> Int
> cost t = sum (zipWith (*) (fringe t) (depths t))

--------------------------------------------------------------------------------
2. Garsia-Wachs algorithm - quadratic version

> gwa :: [Weight] -> Tree Weight
> gwa xs = rebuild xs (build xs)

> rebuild :: [Weight] -> Tree Label -> Tree Weight
> rebuild xs = reduce . zip (map Leaf xs) . sortDepths

> reduce :: [(Tree Label,Depth)] -> Tree Label
> reduce = extract . foldl step []
>  where
>  extract [(t,_)] = t
>  step [] y     = [y]
>  step (x:xs) y = if depth x == depth y
>                  then step xs (join x y)
>                  else y:x:xs
>  join (t1,d) (t2,_) = (Fork t1 t2,d-1)

> depth :: (Tree Label,Depth) -> Depth
> depth = snd

> sortDepths :: Tree Label -> [Depth]
> sortDepths t = elems (array (1,size t) (zip (fringe t) (depths t)))

> build :: [Weight] -> Tree Label
> build ws = extract (foldr combine [] (zip (map Leaf [0..]) (infinity:ws)))
>            where extract [_,(t,_)] = t
>                  infinity = sum ws

> type Pair = (Tree Label,Weight)

> weight :: Pair -> Weight
> weight = snd

> combine :: Pair -> [Pair] -> [Pair]
> combine x (y:z:xs) = if weight x >= weight z 
>                      then combine x (insert (fork y z) xs)
>                      else x:y:z:xs
> combine x xs = x:xs

> fork :: Pair -> Pair -> Pair
> fork (t1,x1) (t2,x2) = (Fork t1 t2,x1+x2)

> insert :: Pair -> [Pair] -> [Pair]
> insert x xs = ys ++ combine x zs where (ys,zs) = split x xs

> split :: Pair -> [Pair] -> ([Pair],[Pair])
> split x xs = span (\y -> weight y < weight x) xs

--------------------------------------------------------------------------------
3. Garsia-Wachs, linearithmic version

First, a revision of build, using the functions

    emptyL  :: List a
    nullL   :: List a -> Bool
    consL   :: a -> List a -> List a
    deconsL :: List a -> (a,List a)
    concatL :: List a -> List a -> List a
    splitL  :: Pair -> List Pair -> (List Pair,List Pair)

> gwaL :: [Weight] -> Tree Weight
> gwaL ws = rebuild ws (buildL ws)

> buildL :: [Weight] -> Tree Label
> buildL ws = extractL (foldr combineL emptyL (zip (map Leaf [0..]) (infinity:ws)))
>             where infinity = sum ws

> extractL :: List Pair -> Tree Label
> extractL xs = t  where ((t,_),_) = deconsL (snd (deconsL xs))

> combineL :: Pair -> List Pair -> List Pair
> combineL x xs = if nullL xs || nullL ys || weight x < weight z
>                 then consL x xs else combineL x (insertL (fork y z) zs)
>                 where (y,ys) = deconsL xs
>                       (z,zs) = deconsL ys

> insertL :: Pair -> List Pair -> List Pair
> insertL x xs = concatL ys (combineL x zs)
>                where (ys,zs) = splitL x xs

> start ws = zip (map Leaf [0..]) (sum ws:ws)

----------------------------------------------------------
4. AVL trees:

> data List a = Null | Node Height (List a) (a,a) (List a)
>               deriving (Eq,Show)
> type Height = Int

> height :: List a -> Height
> height Null           = 0
> height (Node h _ _ _) = h

> node :: List a -> (a, a) -> List a -> List a
> node t1 x t2 = Node h t1 x t2
>    where h = 1 + max (height t1) (height t2)


> balance :: List a -> (a, a) -> List a -> List a
> balance t1 x t2
>     | abs (h1-h2)<=1   = node     t1 x t2
>     | h1==h2+2         = rotateR  t1 x t2
>     | h2==h1+2         = rotateL  t1 x t2 
>     where h1 = height t1; h2 = height t2

> rotateR t1 x t2 = if 0 <= bias t1 then rotr (node t1 x t2)
>                                   else rotr (node (rotl t1) x t2)
> rotateL t1 x t2 = if bias t2 <= 0 then rotl (node t1 x t2) 
>                                   else rotl (node t1 x (rotr t2))

> bias :: List a -> Int
> bias (Node _ t1 x t2) = height t1 - height t2

> rotr (Node _ (Node _ t1 y t2) x t3) = node t1 y (node t2 x t3)
> rotl (Node _ t1 y (Node _ t2 z t3)) = node (node t1 y t2) z t3

> gbalance :: List a -> (a, a) -> List a -> List a
> gbalance t1 x t2
>      | abs (h1-h2) <= 2 = balance  t1 x t2
>      | h1 > h2+2        = balanceR t1 x t2
>      | h2 > h1+2        = balanceL t1 x t2
>      where h1 = height t1; h2 = height t2

> balanceR (Node _ l y r) x t2 = if height r >= height t2 + 2
>                                then balance l y (balanceR r x t2)
>                                else balance l y (node r x t2)
> balanceL t1 x (Node _ l y r) = if height l >= height t1 + 2
>                                then balance (balanceL t1 x l) y r
>                                else balance (node t1 x l) y r

-------------------------------------------------------------
5. Implementation of _L functions:

> emptyL :: List a
> emptyL = Null

> nullL :: List a -> Bool
> nullL Null = True
> nullL _    = False

> consL :: a -> List a -> List a
> consL x Null = node Null (x,x) Null
> consL x (Node _ t1 (y,z) t2) = if nullL t1 
>                                then balance (consL x t1) (x,z) t2 
>                                else balance (consL x t1) (y,z) t2

> deconsL :: List a -> (a,List a)
> deconsL (Node _ t1 xy t2) = if nullL t1 then (snd xy,t2)
>                             else (y, balance t3 xy t2)
>                             where (y,t3) = deconsL t1

> concatL :: List a -> List a -> List a
> concatL t1 Null = t1
> concatL Null t2 = t2
> concatL t1 t2   = gbalance t1 (lastL t1,y) t3
>                   where (y,t3) = deconsL t2

> lastL :: List a -> a
> lastL (Node _ _ xy r) = if nullL r then snd xy else lastL r

> data Piece a = LP (List a) (a,a) | RP (a,a) (List a)

> splitL :: Pair -> List Pair -> (List Pair,List Pair)
> splitL x t = sew (pieces x t)

> pieces :: Pair -> List Pair -> [Piece Pair]
> pieces x t = addPieces t [] where
>   addPieces Null ps = ps
>   addPieces (Node _ t1 (y,z) t2) ps =
>      if w > max (weight y) (weight z)
>      then addPieces t2 (LP t1 (y,z):ps)
>      else addPieces t1 (RP (y,z) t2:ps)
>   w = weight x
 
> sew :: [Piece Pair] -> (List Pair,List Pair)
> sew = foldl step (Null,Null)
>  where step (t1,t2) (LP t x)  = (gbalance t x t1,t2)
>        step (t1,t2) (RP x t)  = (t1,gbalance t2 x t)

-------------------------------------------------------
6. A worst case:

> worst :: Int -> [Int]
> worst n = take (2*n) (twist n)
> twist k = [k,k] ++ twist (k+1)