Huffman.hs

--------------------------------------------------------------------------------
-- Huffman coding
--------------------------------------------------------------------------------

import Data.Maybe
import Data.List (nub, sort)


data Huffman a = Tip Int a
               | Node Int (Huffman a) (Huffman a)
               deriving (Show, Eq, Ord)


data Step = L
          | R
          deriving (Show, Eq)


type Path = [Step]


insert :: (Ord a) => a -> [a] -> [a]
insert y [] = [y]
insert y xx@(x:xs)
  | y < x     = y : xx
  | otherwise = x : insert y xs


count :: (Eq a) => a -> [a] -> (a, Int)
count x xs
  = (x, num)
  where
    num = foldr (\y a -> if x == y then a + 1 else a) 0 xs


countAll :: (Eq a) => [a] -> [a] -> [(a, Int)]
countAll xs ys
  = map (\x -> count x xs) ys


table :: (Eq a) => [a] -> [(a, Int)]
table xs
  = nub $ countAll xs xs


merge :: Huffman a -> Huffman a -> Huffman a
merge ha hb
  | countA < countB = Node sum ha hb
  | otherwise       = Node sum hb ha
  where
    countA = count ha
    countB = count hb
    sum = countA + countB

    count :: Huffman a -> Int
    count (Tip x c) = x
    count (Node x l r) = x


reduce :: (Ord a) => [Huffman a] -> Huffman a
reduce [ha]
  = ha
reduce (ha : hb : hs)
  = reduce $ (merge ha hb) : hs


buildTree :: (Eq a, Ord a) => [a] -> Huffman a
buildTree xs
  = reduce $ sort $ map (\(c, x) -> Tip x c) $ table xs


encode :: (Eq a, Ord a) => [a] -> Huffman a -> Path
encode xs t
  = foldr (\x y -> (fromJust $ code x t) ++ y) [] xs
  where
    code :: (Eq a) => a -> Huffman a -> Maybe [Step]
    code x (Tip _ c)
      | x == c          = Just ([] :: [Step])
      | otherwise       = Nothing
    code x (Node _ lt rt)
      | left /= Nothing  = Just $ L : fromJust left
      | right /= Nothing = Just $ R : fromJust right
      where
        left  = code x lt
        right = code x rt
    code _ _ = Nothing


decode :: (Eq a) => Path -> Huffman a -> [a]
decode [] _ = []
decode xs t
  = ch : decode rem t
  where
    (ch, rem) = code xs t
    
    code :: Path -> Huffman a -> (a, Path)
    code xs (Tip _ c) = (c, xs)
    code (x:xs) (Node _ lt rt)
      | x == L = code xs lt
      | x == R = code xs rt