ハフマン符号化実装メモ

huffman.rs

use std::collections::HashMap;
use std::collections::BinaryHeap;

fn main() {
    let input_bytes = include_bytes!("/path/to/file");
    let mut count = HashMap::new();
    for b in input_bytes.iter() {
        *count.entry(*b).or_insert(0) += 1;
    }
    println!("{:?}", count);

    let mut heap = BinaryHeap::new();
    for (k, v) in count {
        heap.push(Node::new(k, v));
    }

    while heap.len() > 1 {
        let a = heap.pop().unwrap();
        let b = heap.pop().unwrap();
        heap.push(a.merge(b));
    }
    let huffman = Huffman { table: heap.pop().unwrap().to_table() };
    for (k, v) in huffman.table.iter().enumerate().filter(|&(_, ref v)| !v.is_empty()) {
        println!("# {}: {}", k, v.len());
    }

    let mut encoded = 0;
    for b in input_bytes.iter() {
        encoded += huffman.table[*b as usize].len();
    }
    println!("{} => {}", input_bytes.len(), encoded / 8);
}

#[derive(Ord,PartialOrd,Eq,PartialEq)]
struct Node {
    count: isize,
    label: Option<u8>,
    left: Option<Box<Node>>,
    right: Option<Box<Node>>,
}
impl Node {
    fn new(label: u8, count: isize) -> Self {
        Node {
            label: Some(label),
            count: -count,
            left: None,
            right: None,
        }
    }
    fn merge(self, other: Self) -> Self {
        let count = self.count + other.count;
        Node {
            count: count,
            label: None,
            left: Some(Box::new(self)),
            right: Some(Box::new(other)),
        }
    }
    fn to_table(self) -> Vec<Vec<bool>> {
        let mut table = vec![Vec::new(); 0x100];
        self.to_table_impl(&mut Vec::new(), &mut table[..]);
        table
    }
    fn to_table_impl(self, bits: &mut Vec<bool>, table: &mut [Vec<bool>]) {
        if let Some(label) = self.label {
            table[label as usize] = bits.clone();
        } else {
            bits.push(false);
            self.left.unwrap().to_table_impl(bits, table);
            bits.pop();
            bits.push(true);
            self.right.unwrap().to_table_impl(bits, table);
            bits.pop();
        }
    }
}

struct Huffman {
    table: Vec<Vec<bool>>,
}