#![allow(unused)]
use std::io::Write;
use std::{
fs::File,
io::{BufRead, BufReader, BufWriter},
};
#[derive(Debug)]
struct HuffmanNonTerminalNode {
left: u32,
right: u32,
}
// Nodes with id < N are leaves, rest are interior
#[derive(Debug)]
struct HuffmanTree {
priorities: Vec<u64>,
non_terminal_nodes: Vec<HuffmanNonTerminalNode>,
encodings: Vec<Encoding>,
vocab_size: usize,
}
#[derive(Debug, Clone, Default)]
struct Encoding {
val: u64,
num_bits: u8,
}
fn populate_encodings(
node_idx: usize,
vocab_size: usize,
priorities: &Vec<u64>,
non_terminal_nodes: &Vec<HuffmanNonTerminalNode>,
encodings: &mut Vec<Encoding>,
current_encoding: Encoding,
) {
if node_idx < vocab_size {
encodings[node_idx] = current_encoding;
} else {
let left_child = non_terminal_nodes[node_idx - vocab_size].left as usize;
let left_encoding = Encoding {
val: current_encoding.val * 2,
num_bits: current_encoding.num_bits + 1,
};
populate_encodings(
left_child,
vocab_size,
priorities,
non_terminal_nodes,
encodings,
left_encoding,
);
let right_child = non_terminal_nodes[node_idx - vocab_size].right as usize;
let right_encoding = Encoding {
val: current_encoding.val * 2 + 1,
num_bits: current_encoding.num_bits + 1,
};
populate_encodings(
right_child,
vocab_size,
priorities,
non_terminal_nodes,
encodings,
right_encoding,
);
}
}
fn create_huffman_tree(leaves: Vec<u64>) -> (HuffmanTree, u64) {
let vocab_size: usize = leaves.len();
let mut priorities = leaves;
let mut encodings: Vec<Encoding> = vec![Encoding::default(); vocab_size];
let mut total_num_bits = 0u64;
let mut original_prio_head = 0usize;
let mut non_terminal_node_head = vocab_size;
let mut non_terminal_nodes: Vec<HuffmanNonTerminalNode> = Vec::new();
while original_prio_head < vocab_size || non_terminal_node_head < (priorities.len() - 1) {
let left = if original_prio_head >= vocab_size {
non_terminal_node_head += 1;
non_terminal_node_head - 1
} else if non_terminal_node_head >= priorities.len() {
original_prio_head += 1;
original_prio_head - 1
} else if priorities[original_prio_head] < priorities[non_terminal_node_head] {
original_prio_head += 1;
original_prio_head - 1
} else {
non_terminal_node_head += 1;
non_terminal_node_head - 1
};
let right = if original_prio_head >= vocab_size {
non_terminal_node_head += 1;
non_terminal_node_head - 1
} else if non_terminal_node_head >= priorities.len() {
original_prio_head += 1;
original_prio_head - 1
} else if priorities[original_prio_head] < priorities[non_terminal_node_head] {
original_prio_head += 1;
original_prio_head - 1
} else {
non_terminal_node_head += 1;
non_terminal_node_head - 1
};
let prio = priorities[left] + priorities[right];
priorities.push(prio);
let left = left as u32;
let right = right as u32;
let new_node = HuffmanNonTerminalNode { left, right };
total_num_bits += prio;
non_terminal_nodes.push(new_node);
}
populate_encodings(
priorities.len() - 1,
vocab_size,
&priorities,
&non_terminal_nodes,
&mut encodings,
Encoding::default(),
);
let tree = HuffmanTree {
priorities,
non_terminal_nodes,
encodings,
vocab_size,
};
(tree, total_num_bits)
}
fn main() {
let leaves = {
let input_file: File = File::open("huffman.in").unwrap();
let mut input_file: BufReader<File> = BufReader::new(input_file);
#[allow(non_snake_case)]
let mut N = String::new();
let _ = input_file.read_line(&mut N).unwrap();
input_file
.lines()
.map(|line| line.unwrap().trim().parse::<u64>().unwrap())
.collect()
};
let (tree, total_num_bits) = create_huffman_tree(leaves);
let mut output_file = BufWriter::new(File::create("huffman.out").unwrap());
let _ = writeln!(output_file, "{total_num_bits}").unwrap();
for encoding in tree.encodings {
let _ = writeln!(output_file, "{} {}", encoding.num_bits, encoding.val);
}
}