huffman_idk/
lib.rs

1// Copyright (C) 2024 Etham Uppal. All rights reserved.
2
3use std::{cmp, collections::BinaryHeap};
4
5use baa::BitVecValue;
6
7struct Letter {
8    index: usize,
9    frequency: f64,
10}
11
12impl PartialEq for Letter {
13    fn eq(&self, other: &Self) -> bool {
14        self.index == other.index && self.frequency == other.frequency
15    }
16}
17
18impl Eq for Letter {}
19
20impl PartialOrd for Letter {
21    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
22        Some(self.cmp(other))
23    }
24}
25
26impl Ord for Letter {
27    fn cmp(&self, other: &Self) -> cmp::Ordering {
28        self.frequency.total_cmp(&other.frequency)
29    }
30}
31
32#[derive(Debug)]
33enum Node {
34    Leaf { value: char },
35    Interior { left: usize, right: usize },
36}
37
38fn construct_prefix_code(
39    result: &mut Vec<(char, BitVecValue)>,
40    tree: &[Node],
41    current: usize,
42    code: u64,
43    depth: u32,
44) {
45    match &tree[current] {
46        Node::Leaf { value } => {
47            result.push((*value, BitVecValue::from_u64(code, depth)));
48        }
49        Node::Interior { left, right } => {
50            #[allow(clippy::identity_op)]
51            construct_prefix_code(
52                result,
53                tree,
54                *left,
55                (code << 1) | 0,
56                depth + 1,
57            );
58            construct_prefix_code(
59                result,
60                tree,
61                *right,
62                (code << 1) | 1,
63                depth + 1,
64            );
65        }
66    }
67}
68
69pub fn compute_prefix_code(
70    alphabet: &[(char, f64)],
71) -> Vec<(char, BitVecValue)> {
72    if alphabet.len() <= 2 {
73        alphabet
74            .iter()
75            .enumerate()
76            .map(|(index, (value, _))| {
77                (*value, BitVecValue::from_u64(index as u64, 1))
78            })
79            .collect()
80    } else {
81        let mut tree = alphabet
82            .iter()
83            .cloned()
84            .map(|(value, _)| Node::Leaf { value })
85            .collect::<Vec<_>>();
86        let mut letters =
87            BinaryHeap::from_iter((0..alphabet.len()).map(|index| {
88                cmp::Reverse(Letter {
89                    index,
90                    frequency: alphabet[index].1,
91                })
92            }));
93
94        while !letters.is_empty() {
95            let cmp::Reverse(first) = letters.pop().expect("invariant");
96            let cmp::Reverse(second) = letters.pop().expect("invariant");
97
98            tree.push(Node::Interior {
99                left: first.index,
100                right: second.index,
101            });
102
103            if !letters.is_empty() {
104                letters.push(cmp::Reverse(Letter {
105                    index: tree.len() - 1,
106                    frequency: first.frequency + second.frequency,
107                }));
108            }
109        }
110
111        let mut result = vec![];
112
113        assert!(!tree.is_empty());
114        construct_prefix_code(&mut result, &tree, tree.len() - 1, 0, 0);
115
116        result
117    }
118}
huffman_idk/lib.rs

huffman_idk/
lib.rs
