Cod sursa(job #3228409)

Utilizator andreioneaAndrei Onea andreionea Data 8 mai 2024 10:21:25
Problema Coduri Huffman Scor 30
Compilator rs Status done
Runda Arhiva educationala Marime 3.84 kb
#![allow(non_snake_case)]
use std::{cell::RefCell, collections::{BinaryHeap, HashMap}, fs::File, io::{BufRead, BufReader, BufWriter}, rc::Rc};
use std::io::Write;

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

#[derive(Debug)]
struct Encoding {
    val: u32,
    num_bits: u32,
}

impl HuffmanNode {
    fn prio(&self) -> u32 {
        match self{
            HuffmanNode::Leaf { prio, chr: _ } => *prio,
            HuffmanNode::Internal { prio, left: _, right: _ } => *prio,
        }
    }
    fn new_leaf(prio: u32, chr: u32) -> Rc<RefCell<Self>> {
        Rc::new(RefCell::new(HuffmanNode::Leaf { prio, chr }))
    }
    fn new_inner(left: Rc<RefCell<Self>>, right: Rc<RefCell<Self>>) -> Rc<RefCell<Self>> {
        let prio = left.borrow().prio() + right.borrow().prio();
        let node: HuffmanNode = HuffmanNode::Internal { prio , left, right };
        Rc::new(RefCell::new(node))
    }
    fn create_encoding_impl(node: &Self, existing_encodings: &mut HashMap<u32, Encoding>, current_encoding: Encoding) {
        match node{
            HuffmanNode::Leaf { prio: _, chr } => {
                existing_encodings.insert(*chr, 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.borrow(), existing_encodings, left_encoding);
                Self::create_encoding_impl(&right.borrow(), existing_encodings, right_encoding);
            },
        }
    }
    fn create_encoding(&self) -> HashMap<u32, Encoding> {
        let mut encodings = HashMap::new();
        let zero_encoding = Encoding{ val: 0, num_bits: 0 };
        Self::create_encoding_impl(self, &mut encodings, zero_encoding);
        encodings
    }
}

impl PartialOrd for HuffmanNode {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        other.prio().partial_cmp(&self.prio())
    }
}

impl Ord for HuffmanNode {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        other.prio().cmp(&self.prio())
    }
}

fn create_huffman_tree(T: &Vec<u32>) -> (Rc<RefCell<HuffmanNode>>, u32) {
    if T.is_empty() {
        panic!("Nopenopenopenope");
    }
    let mut h: BinaryHeap<Rc<RefCell<HuffmanNode>>> = T.iter().enumerate().map(|(idx, val)| {
        HuffmanNode::new_leaf(*val, idx as u32)
    }).collect();
    let mut total_num_bits = 0;
    while h.len() >= 2 {
        let left = h.pop().unwrap();
        let right = h.pop().unwrap();
        let new_node = HuffmanNode::new_inner(left, right);
        total_num_bits += new_node.borrow().prio();
        h.push(new_node);
    }
    (h.pop().unwrap(), total_num_bits)
}

fn main() {
    let input_file = File::open("huffman.in").unwrap();
    let mut input_file = BufReader::new(input_file);
    let mut N = String::new();
    let _ = input_file.read_line(&mut N).unwrap();
    let N = N.trim().parse::<u32>().unwrap();
    let T: Vec<u32> = input_file.lines().map(|line| line.unwrap().trim().parse::<u32>().unwrap()).collect();
    let (H, total_num_bits) = create_huffman_tree(&T);
    let encodings = H.borrow().create_encoding();
    let mut output_file = BufWriter::new(File::create("huffman.out").unwrap());
    let _ = writeln!(output_file, "{total_num_bits}").unwrap();
    for i in 0..N {
        let _ = writeln!(output_file, "{} {}", encodings.get(&i).unwrap().num_bits, encodings.get(&i).unwrap().val).unwrap();
    }
}