Huffman.hs

module Huffman where

import Data.Map
import Data.Word
import Data.Ord
import Data.List
import Data.Maybe

data Tree a = Node {small::Tree a, big::Tree a}
            | Leaf {occurrences::Int, symbol::a}
            deriving (Show, Eq)

instance Ord a => Ord (Tree a) where
  compare = comparing frequency

frequency :: Tree a -> Int
frequency l@(Leaf _ _) = occurrences l
frequency n@(Node _ _) = (frequency . small $ n) + (frequency . big $ n)

initialTrees :: Ord a => [a] -> [Tree a]
initialTrees = sort . Data.List.map convertToLeaf . toList . weights
  where convertToLeaf (s,f) = Leaf { occurrences=f, symbol=s }
        weights = foldl update empty
          where update acc symbol = insertWith (+) symbol 1 acc

codeTree :: Ord a => [Tree a] -> Tree a
codeTree (tree:[]) = tree
codeTree (a:b:trees) = codeTree $ Data.List.insert (Node a b) trees

trav :: Ord a => Tree a -> Map a String
trav (Node a b) = fromList $ trav' a "0" ++ trav' b "1"
  where trav' (Node a b) path = trav' a (path ++ "0") ++ trav' b (path ++ "1")
        trav' (Leaf _occurrences symbol) path = [(symbol, path)]

codex :: Ord k => [k] -> Map k String
codex = trav . codeTree . initialTrees

encode :: Ord k => Map k String -> [k] -> String
encode table symbols = concat [ table ! s| s <- symbols ]

decode :: Ord k => Map String k -> String -> [k]
decode table [] = []
decode table bits = decode' table bits 1
  where decode' table bits n = let (prefix, rest) = splitAt n bits
                               in case Data.Map.lookup prefix table of
                                    Just s -> s : decode table rest
                                    Nothing -> decode' table bits (n + 1)

s = "qazwsxqazqaz"
etable = codex s
dtable = fromList [ (v,k) | (k,v) <- toList . codex $ s ]
bits = encode etable s