Cod sursa(job #3228525)

Utilizator andreioneaAndrei Onea andreionea Data 8 mai 2024 17:11:02
Problema Coduri Huffman Scor 100
Compilator rs Status done
Runda Arhiva educationala Marime 4.58 kb
#![allow(non_snake_case)]
use std::io::Write;
use std::{
    fs::File,
    io::{BufRead, BufReader, BufWriter},
};


#[derive(Debug)]
struct HuffmanNode {
    prio: u64,
    left: u32,
    right: u32,
}

type HuffmanTree = Vec<HuffmanNode>;

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

impl HuffmanNode {
    fn new_leaf(prio: u64, chr: u32) -> Self {
        HuffmanNode {
            prio,
            left: chr,
            right: 0
        }
    }
    fn new_inner(tree: &HuffmanTree, left: u32, right: u32) -> Self {
        let prio = tree[left as usize].prio + tree[right as usize].prio;
        HuffmanNode {
            prio,
            left,
            right
        }
    }
    fn create_encoding_impl(
        tree: &HuffmanTree,
        node_idx: u32,
        N: u32,
        existing_encodings: &mut Vec<Encoding>,
        current_encoding: Encoding,
    ) {
        if node_idx < N {
            existing_encodings[node_idx as usize] = current_encoding;
        } else {
            let left = tree[node_idx as usize].left;
            let right = tree[node_idx as usize].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(
                tree,
                left,
                N,
                existing_encodings,
                left_encoding,
            );
            Self::create_encoding_impl(
                tree,
                right,
                N,
                existing_encodings,
                right_encoding,
            );
        }
    }
    fn create_encoding(&self, tree: &HuffmanTree, 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(tree, (tree.len() - 1) as u32, N as u32, &mut encodings, zero_encoding);
        encodings
    }
}

fn create_huffman_tree(mut tree: HuffmanTree) -> (HuffmanTree, u64) {
    let mut queue1_head = 0usize;
    let queue1_tail = tree.len();
    let mut queue2_head = tree.len();

    let mut total_num_bits = 0u64;
    while queue1_head < queue1_tail || queue2_head < (tree.len() - 1) {
        let left = if queue1_head >= queue1_tail {
            queue2_head += 1;
            queue2_head - 1
        } else if queue2_head >= tree.len() {
            queue1_head += 1;
            queue1_head - 1
        } else if tree[queue1_head].prio < tree[queue2_head].prio {
            queue1_head += 1;
            queue1_head - 1
        } else {
            queue2_head += 1;
            queue2_head - 1
        };

        let right = if queue1_head >= queue1_tail {
            queue2_head += 1;
            queue2_head - 1
        } else if queue2_head >= tree.len() {
            queue1_head += 1;
            queue1_head - 1
        } else if tree[queue1_head].prio < tree[queue2_head].prio {
            queue1_head += 1;
            queue1_head - 1
        } else {
            queue2_head += 1;
            queue2_head - 1
        };

        let new_node = HuffmanNode::new_inner(&tree, left as u32, right as u32);
        total_num_bits += new_node.prio;
        tree.push(new_node);
    }
    (tree, total_num_bits)
}

fn main() {
    let (tree, N) = {
        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();
        (
            input_file
                .lines()
                .enumerate()
                .map(|(idx, line)| {
                    let vi = line.unwrap().trim().parse::<u64>().unwrap();
                    HuffmanNode::new_leaf(vi, idx as u32)
                })
                .collect(),
            N,
        )
    };
    let (tree, total_num_bits) = create_huffman_tree(tree);
    let root = tree.last().unwrap();
    let encodings = root.create_encoding(&tree, 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);
    }
}