Huffman.hs

module Huffman where

import Data.List
import Data.SortedList

-- A symbol and its probability
type Symbol = String
data SymbolP = SymbolP { symbol :: String
                       , probaS :: Float
                       } deriving (Show)
instance Eq SymbolP where
  SymbolP _ p1 == SymbolP _ p2 = p1 == p2
instance Ord SymbolP where
  (SymbolP _ p1) `compare` (SymbolP _ p2) = p1 `compare` p2

-- Level 0: list of SymbolP
-- [("a", 0.3), ("b", 0.3), ("c", 0.2), ("d", 0.2)]
-- Level 1: 1 simple composite
-- [("a", 0.3), ("b", 0.3), COMPOSITE (("c", 0.2), ("d", 0.2)) 0.4]
-- Level 2: nested composite
-- [("a", 0.3), COMPOSITE ((COMPOSITE (("c", 0.2), ("d", 0.2)) 0.4), ("b", 0.3)) 0.7]

data HTree = HTree { left :: HNode, right :: HNode, probaT :: Float } deriving (Show)
instance Eq HTree where
  HTree _ _ t1 == HTree _ _ t2 = t1 == t2
instance Ord HTree where
  (HTree _ _ t1) `compare` (HTree _ _ t2) = t1 `compare` t2
data HNode = Leaf SymbolP | SubTree HTree deriving (Show)
instance Eq HNode where
  hn1 == hn2 = getProba hn1 == getProba hn2
instance Ord HNode where
  hn1 `compare` hn2 = compare (getProba hn1) (getProba hn2)

data HCode = HCode { symbolP :: SymbolP
                   , code :: String
                   } deriving (Show)

getProba :: HNode -> Float
getProba (Leaf s) = probaS s
getProba (SubTree t) = probaT t

-- The 2 first elements have the smallest probability, collapses them
collapseTwoSmallest :: SortedList HNode -> SortedList HNode
collapseTwoSmallest nodes =
  Data.SortedList.insert collapsed rest
    where rest = Data.SortedList.drop 2 nodes
          smallest = Data.SortedList.fromSortedList $ Data.SortedList.take 2 nodes
          leftElem = Data.List.head smallest
          rightElem = Data.List.last smallest
          collapsed = SubTree $ HTree leftElem rightElem (getProba leftElem + getProba rightElem)

buildHuffmanTree :: SortedList HNode -> HNode
buildHuffmanTree nodes
  | null (fromSortedList nodes) = Leaf $ SymbolP "" 0
  | Data.List.length (fromSortedList nodes) == 1 = Data.List.head $ fromSortedList nodes
  | otherwise = buildHuffmanTree $ collapseTwoSmallest nodes

buildCodes :: String -> HNode -> [HCode]
buildCodes prefix (Leaf a) = [HCode a prefix]
buildCodes prefix (SubTree tree) =
  buildCodes (prefix ++ "0") (left tree) ++ buildCodes (prefix ++ "1") (right tree)

omg, I find that so complicated.... I guess I'm missing practice