Huffman example.

2 files changed, 112 insertions, 0 deletions

diff --git a/rust/huffman/Cargo.toml b/rust/huffman/Cargo.toml
new file mode 100644
index 0000000..313fad4
--- /dev/null
+++ b/rust/huffman/Cargo.toml
@@ -0,0 +1,7 @@
+[package]
+name = "huffman"
+version = "0.1.0"
+authors = ["Kjetil Orbekk <kjetil.orbekk@gmail.com>"]
+
+[dependencies]
+itertools = "*"
diff --git a/rust/huffman/src/main.rs b/rust/huffman/src/main.rs
new file mode 100644
index 0000000..d2fc8a1
--- /dev/null
+++ b/rust/huffman/src/main.rs
@@ -0,0 +1,105 @@
+extern crate itertools;
+use itertools::Itertools;
+use std::cmp::Eq;
+use std::cmp::Ordering;
+use std::collections::binary_heap::BinaryHeap;
+
+static MESSAGE: &'static str = r#"
+  Hello! What's going on? This is a program.
+  The brown fox jumped over the lazy dog.
+  xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
+  \o/
+"#;
+
+#[derive(Debug, Eq, PartialEq)]
+enum Tree<T> {
+    Leaf(i32, T),
+    Branch(i32, Box<Tree<T>>, Box<Tree<T>>),
+}
+
+#[derive(Debug, Eq, PartialEq)]
+enum Direction { Left, Right }
+
+fn cost<T>(t: &Tree<T>) -> i32 {
+    match *t {
+        Tree::Leaf(cost, _) => cost,
+        Tree::Branch(cost, _, _) => cost,
+    }
+}
+
+impl<T: Eq> Ord for Tree<T> {
+    fn cmp(&self, other: &Self) -> Ordering {
+        cost(other).cmp(&cost(self))
+    }
+}
+
+impl<T: Eq> PartialOrd for Tree<T> {
+    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+fn make_path<T: Eq>(tree: &Tree<T>, elem: &T) -> Vec<Direction> {
+    fn find<T: Eq>(n: &Tree<T>, elem: &T) -> Option<Vec<Direction>> {
+        match *n {
+            Tree::Leaf(_, ref e) =>
+                if e == elem { Some(vec!()) } else { None },
+            Tree::Branch(_, ref left, ref right) =>
+                find(left.as_ref(), elem)
+                .map(|mut path| { path.push(Direction::Left); path })
+
+                .or(find(right.as_ref(), elem)
+                    .map(|mut path| { path.push(Direction::Right); path})),
+        }
+    };
+    let mut result = find(tree, elem).unwrap();
+    result.reverse();
+    result
+}
+
+fn find_char<T: Eq+Copy>(tree: &Tree<T>, path: &[Direction]) -> Option<T> {
+    match (tree, path.split_first()) {
+        (&Tree::Branch(_, ref left, _), Some((&Direction::Left, ds))) =>
+            find_char(left.as_ref(), ds),
+        (&Tree::Branch(_, _, ref right), Some((&Direction::Right, ds))) =>
+            find_char(right.as_ref(), ds),
+        (&Tree::Leaf(_, ref c), None) => Some(*c),
+        _ => None
+    }
+}
+
+fn main() {
+    if MESSAGE.is_empty() {
+        panic!("Empty string");
+    }
+
+    let mut counts = MESSAGE
+        .chars()
+        .sorted().into_iter()
+        .group_by(|x| *x)
+        .map(|(elem, xs)| (elem, xs.len() as i32))
+        .map(|(elem, cost)| Box::new(Tree::Leaf(cost, elem)))
+        .collect::<BinaryHeap<_>>();
+
+    let mut left;
+    loop {
+        left = counts.pop().unwrap();
+        let right = match counts.pop() {
+            Some(right) => right,
+            None => break
+        };
+
+        counts.push(Box::new(
+            Tree::Branch(cost(left.as_ref()) + cost(right.as_ref()),
+                         left, right)));
+    }
+
+    let restored_message = MESSAGE
+        .chars()
+        .map(|c| make_path(left.as_ref(), &c))
+        .flat_map(|p| find_char(left.as_ref(), &p))
+        .collect::<String>();
+
+    println!("{}", MESSAGE);
+    println!("{}", restored_message);
+}