struct HuffmanNode {
total_count: u32,
index_left: i16,
index_right_or_value: i16, }
const SENTINEL: HuffmanNode = HuffmanNode {
total_count: u32::MAX,
index_left: -1,
index_right_or_value: -1,
};
fn set_depth(idx: usize, pool: &[HuffmanNode], depth: &mut [u8], level: u8) {
let node = &pool[idx];
if node.index_left >= 0 {
let new_level = level + 1;
set_depth(node.index_left as usize, pool, depth, new_level);
set_depth(node.index_right_or_value as usize, pool, depth, new_level);
} else {
depth[node.index_right_or_value as usize] = level;
}
}
pub(crate) fn create_huffman_tree(data: &[u32], tree_limit: u8, depth: &mut [u8]) {
let length = data.len();
debug_assert!(depth.len() >= length);
for d in depth[..length].iter_mut() {
*d = 0;
}
let mut count_limit: u32 = 1;
loop {
let mut tree: Vec<HuffmanNode> = Vec::with_capacity(2 * length + 1);
for (i, &data) in data[..length].iter().enumerate().rev() {
if data != 0 {
let count = data.max(count_limit.saturating_sub(1));
tree.push(HuffmanNode {
total_count: count,
index_left: -1,
index_right_or_value: i as i16,
});
}
}
let n = tree.len();
if n == 1 {
depth[tree[0].index_right_or_value as usize] = 1;
return;
}
tree.sort_by_key(|a| a.total_count);
tree.push(SENTINEL);
tree.push(SENTINEL);
let mut i = 0usize; let mut j = n + 1;
for _ in 0..(n - 1) {
let left;
if tree[i].total_count <= tree[j].total_count {
left = i;
i += 1;
} else {
left = j;
j += 1;
}
let right;
if tree[i].total_count <= tree[j].total_count {
right = i;
i += 1;
} else {
right = j;
j += 1;
}
let parent = tree.len() - 1;
tree[parent].total_count = tree[left].total_count + tree[right].total_count;
tree[parent].index_left = left as i16;
tree[parent].index_right_or_value = right as i16;
tree.push(SENTINEL);
}
debug_assert_eq!(tree.len(), 2 * n + 1);
set_depth(2 * n - 1, &tree, depth, 0);
let max_depth = depth[..length].iter().copied().max().unwrap_or(0);
if max_depth <= tree_limit {
return;
}
count_limit = count_limit
.checked_mul(2)
.expect("huffman depth limit unreachable");
}
}