Cod sursa(job #3228579)

Utilizator andreioneaAndrei Onea andreionea Data 8 mai 2024 23:05:03
Problema Heapuri cu reuniune Scor 0
Compilator rs Status done
Runda Arhiva educationala Marime 4.56 kb
#![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);
    }
}