Cod sursa(job #3228461)

Utilizator andreioneaAndrei Onea andreionea Data 8 mai 2024 12:31:30
Problema Coduri Huffman Scor 70
Compilator rs Status done
Runda Arhiva educationala Marime 4.74 kb
#![allow(non_snake_case)]
use std::{collections::VecDeque, fs::File, io::{BufRead, BufReader, BufWriter}};
use std::io::Write;

#[derive(PartialEq, Eq, Debug)]
enum HuffmanNode {
    Leaf {
        prio: u64,
        chr: u32
    },
    Internal {
        prio: u64,
        left: Box<HuffmanNode>,
        right: Box<HuffmanNode>,
    }
}

#[derive(Debug, Clone)]
struct Encoding {
    val: u64,
    num_bits: u8,
}

impl HuffmanNode {
    fn prio(&self) -> u64 {
        match self{
            HuffmanNode::Leaf { prio, chr: _ } => *prio,
            HuffmanNode::Internal { prio, left: _, right: _ } => *prio,
        }
    }
    fn new_leaf(prio: u64, chr: u32) -> Self {
        HuffmanNode::Leaf { prio, chr }
    }
    fn new_inner(left: Self, right: Self) -> Self {
        let prio = left.prio() + right.prio();
        HuffmanNode::Internal { prio , left: Box::new(left), right: Box::new(right) }
    }
    fn create_encoding_impl(node: &Self, existing_encodings: &mut Vec<Encoding>, current_encoding: Encoding) {
        match node{
            HuffmanNode::Leaf { prio: _, chr } => {
                existing_encodings[*chr as usize] = current_encoding;
            },
            HuffmanNode::Internal { prio: _, left, right } => {
                let left_encoding = Encoding {
                    val: current_encoding.val * 2,
                    num_bits: current_encoding.num_bits + 1,
                };
                let right_encoding = Encoding {
                    val: current_encoding.val * 2 + 1,
                    num_bits: current_encoding.num_bits + 1,
                };
                Self::create_encoding_impl(left.as_ref(), existing_encodings, left_encoding);
                Self::create_encoding_impl(right.as_ref(), existing_encodings, right_encoding);
            },
        }
    }
    fn create_encoding(&self, N: usize) -> Vec<Encoding> {
        let mut encodings = vec![Encoding{ val: 0, num_bits: 0 };N];
        let zero_encoding = Encoding{ val: 0, num_bits: 0 };
        Self::create_encoding_impl(self, &mut encodings, zero_encoding);
        encodings
    }
}

fn create_huffman_tree(mut new_node_iterator: impl Iterator<Item = HuffmanNode>, N: usize) -> (HuffmanNode, u64) {
    let mut q = VecDeque::with_capacity(2 + N / 2);
    let mut total_num_bits = 0u64;
    let mut last_read_node: Option<HuffmanNode> = None;
    loop {
        match (q.front(), last_read_node.take()) {
            (None, None) => {
                q.push_back(new_node_iterator.next().unwrap())
            },
            (None, Some(val)) => {
                q.push_back(val);
            },
            (Some(_), None) => {
                if let Some(new_node) = new_node_iterator.next() {
                    last_read_node = Some(new_node);
                } else {
                    let left = q.pop_front().unwrap();
                    if q.is_empty() {
                        return (left, total_num_bits);
                    }
                    let right = q.pop_front().unwrap();
                    total_num_bits += left.prio() + right.prio();
                    let new_node = HuffmanNode::new_inner(left, right);
                    q.push_back(new_node);
                }
            },
            (Some(front), Some(last_read)) => {
                if last_read.prio() < front.prio() {
                    q.push_front(last_read);
                } else {
                    let left = q.pop_front().unwrap();
                    let right = if q.is_empty() || q.front().unwrap().prio() > last_read.prio() {
                        last_read
                    } else {
                        last_read_node = Some(last_read);
                        q.pop_front().unwrap()
                    };
                    let new_node = HuffmanNode::new_inner(left, right);
                    total_num_bits += new_node.prio();
                    q.push_back(new_node);
                }
                
            }
        }
    }

}

fn main() {
    let input_file: File = File::open("huffman.in").unwrap();
    let mut input_file: BufReader<File> = BufReader::new(input_file);
    let mut N = String::new();
    let _ = input_file.read_line(&mut N).unwrap();
    let N = N.trim().parse::<usize>().unwrap();
    let new_nodes_iterator = input_file.lines().enumerate().map(|(idx,line)| {
        let vi = line.unwrap().trim().parse::<u64>().unwrap();
        HuffmanNode::new_leaf(vi, idx as u32)
    });
    let (H, total_num_bits) = create_huffman_tree(new_nodes_iterator, N);
    let encodings = H.create_encoding(N);
    let mut output_file = BufWriter::new(File::create("huffman.out").unwrap());
    let _ = writeln!(output_file, "{total_num_bits}").unwrap();
    for encoding in encodings {
        let _ = writeln!(output_file, "{} {}", encoding.num_bits, encoding.val);
    }
}