huffman.rs

use std::cmp::Ordering;
use std::collections::BinaryHeap;

struct Node {
    weight: i32,
    item: Item,
}

impl Ord for Node {
    fn cmp(&self, other: &Self) -> Ordering {
        other.weight.cmp(&self.weight)
    }
}

impl PartialOrd for Node {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl PartialEq for Node {
    fn eq(&self, other: &Self) -> bool {
        self.weight == other.weight
    }
}

impl Eq for Node {}

enum Item {
    Branch(Box<Node>, Box<Node>),
    Leaf(char),
}

fn build_table_impl(node: &Node, table: &mut Vec<(char, String)>, prefix: &mut String) {
    match node.item {
        Item::Branch(ref left, ref right) => {
            prefix.push('0');
            build_table_impl(&left, table, prefix);
            prefix.pop();

            prefix.push('1');
            build_table_impl(&right, table, prefix);
            prefix.pop();
        }
        Item::Leaf(sym) => table.push((sym, prefix.to_owned())),
    };
}

fn build_table(node: &Node) -> Vec<(char, String)> {
    let mut table = Vec::new();
    build_table_impl(&node, &mut table, &mut String::new());
    table
}

fn build_tree(freqs: &[(char, i32)]) -> Node {
    let mut heap = BinaryHeap::new();
    for &(sym, freq) in freqs {
        heap.push(Node {
            weight: freq,
            item: Item::Leaf(sym),
        });
    }

    while heap.len() > 1 {
        let item1 = heap.pop().unwrap();
        let item2 = heap.pop().unwrap();
        heap.push(Node {
            weight: item1.weight + item2.weight,
            item: Item::Branch(Box::new(item1), Box::new(item2)),
        });
    }

    heap.pop().unwrap()
}

fn huffman(freqs: &[(char, i32)]) -> Vec<(char, String)> {
    let tree = build_tree(&freqs);

    let mut table = build_table(&tree);
    table.sort();

    table
}

fn main() {
    // https://en.wikipedia.org/wiki/Letter_frequency
    let freqs = [('a', 8167),
                 ('b', 1492),
                 ('c', 2782),
                 ('d', 4253),
                 ('e', 12702),
                 ('f', 2228),
                 ('g', 2015),
                 ('h', 6094),
                 ('i', 6966),
                 ('j', 153),
                 ('k', 772),
                 ('l', 4025),
                 ('m', 2406),
                 ('n', 6749),
                 ('o', 7507),
                 ('p', 1929),
                 ('q', 95),
                 ('r', 5987),
                 ('s', 6327),
                 ('t', 9056),
                 ('u', 2758),
                 ('v', 978),
                 ('w', 2361),
                 ('x', 150),
                 ('y', 1974),
                 ('z', 74)];

    for &(sym, ref str) in &huffman(&freqs) {
        println!("{}: {}", sym, str);
    }
}