main.rs

extern crate common;
use common::measure_and_print;

use std::collections::{BinaryHeap, HashMap};
use std::io;

#[derive(Debug, Eq, PartialEq)]
enum NodeKind {
    Leaf(u8),
    Branch(Box<Node>, Box<Node>),
}

#[derive(Debug, Eq, PartialEq)]
struct Node {
    frequency: usize,
    kind: NodeKind,
}

impl Ord for Node {
    fn cmp(&self, rhs: &Self) -> std::cmp::Ordering {
        rhs.frequency.cmp(&self.frequency)
    }
}

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

impl Node {
    // create leaf-node with particular byte
    fn leaf(frequency: usize, byte: u8) -> Self {
        Node { frequency, kind: NodeKind::Leaf(byte) }
    }

    // create new node on top of `left` and `right`
    fn branch(left: Self, right: Self) -> Self {
        Node { frequency: left.frequency + right.frequency, kind: NodeKind::Branch(Box::new(left), Box::new(right)) }
    }

    // traverse tree building letter->code `map`
    fn build_map(&self, map: &mut HashMap<u8, String>, prefix: String) {
        match self.kind {
            NodeKind::Branch(ref left_child, ref right_child) => {
                left_child.build_map(map, prefix.clone() + "0");
                right_child.build_map(map, prefix + "1");
            }
            NodeKind::Leaf(byte) => {
                map.insert(byte, prefix);
            }
        }
    }
}

fn calculate_letter_frequencies(string: &str) -> HashMap<u8, usize> {
    let mut frequencies = vec![0usize; 256];
    for byte in string.as_bytes() {
        frequencies[*byte as usize] += 1;
    }

    let mut map = HashMap::<u8, usize>::with_capacity(256);
    for (idx, frequency) in frequencies.into_iter().enumerate() {
        if frequency != 0 {
            map.insert(idx as u8, frequency);
        }
    }
    map
}

pub struct Encoder<'a> {
    m: [&'a str; 256],
}

impl<'a> Encoder<'a> {
    pub fn new(table: &'a mut HashMap<u8, String>) -> Self {
        let mut m = [""; 256];
        for i in 0..=255 {
            if table.contains_key(&i) {
                m[i as usize] = table[&i].as_ref();
            }
        }

        Encoder { m: m }
    }
    fn encode_byte(&self, byte: u8) -> &str {
        self.m[byte as usize]
    }
    pub fn encode(&self, string: &str) -> String {
        let mut out = String::with_capacity(string.len());
        string.as_bytes().iter().for_each(|byte| {
            out.push_str( self.encode_byte(*byte) );
        });
        out
    }
}


fn read_line() -> String {
    let mut line = String::new();
    io::stdin().read_line(&mut line).unwrap();
    line.trim().to_string()
}

fn main() {
    // 1.Read message
    let sentence = read_line();

    let mut letter_freq = HashMap::<u8, usize>::new();
    let mut code_string = String::new();
    let mut table = HashMap::<u8, String>::new();

    measure_and_print(||
        {
            // 2. Calculate frequancy map for letters
            letter_freq = calculate_letter_frequencies(&sentence);

            // 3. Build Huffman tree & table
            let mut heap = BinaryHeap::new();
            for (byte, frequency) in &letter_freq {
                heap.push(Node::leaf(*frequency, *byte));
            }
            if heap.len() > 1 {
                while heap.len() > 1 {
                    // take two least frequent nodes and build new one on top of it
                    let a = heap.pop().expect("heap.pop cannot fail since its len >= 2");
                    let b = heap.pop().expect("heap.pop cannot fail since its len >= 2");
                    heap.push(Node::branch(a, b));
                }
                // take head of the tree and traverse it build letter->code map
                let head = heap.peek().expect("heap should not be empty");
                head.build_map(&mut table, String::from(""));
            } else if heap.len() == 1 {
                // take head of the tree and traverse it build letter->code map
                let head = heap.peek().expect("heap should not be empty");
                head.build_map(&mut table, String::from("0"));
            } else {
                panic!("go check heap yourself")
            }

            // 4. Encode message
            let encoder = Encoder::new(&mut table);
            code_string = encoder.encode(&sentence);
        });

    // 5. Output
    println!("{} {}", letter_freq.len(), code_string.len());
    //for (letter, code) in &table {
    //    println!("{}: {}", letter, code);
    //}
    println!("{}", code_string);
}