Cod sursa(job #3228584)

Utilizator andreioneaAndrei Onea andreionea Data 8 mai 2024 23:13:38
Problema Heapuri cu reuniune Scor 30
Compilator rs Status done
Runda Arhiva educationala Marime 3.43 kb
use std::{collections::BinaryHeap, fs::File, io::{BufRead, BufReader, BufWriter, Write}};

#[derive(Debug, Clone)]
struct HeapFragment(BinaryHeap<u32>);

#[derive(Debug, Clone)]
struct HeapHeap {
    heaps: BinaryHeap<HeapFragment>
}

impl<> PartialEq for HeapFragment {
    fn eq(&self, other: &Self) -> bool {
        if self.0.is_empty() != other.0.is_empty() {
            false
        } else if self.0.is_empty() {
            true
        } else {
            self.0.peek().unwrap() == other.0.peek().unwrap()
        }
    }
}

impl<> Eq for HeapFragment {

}

impl<> PartialOrd for HeapFragment {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        if self == other {
            Some(std::cmp::Ordering::Equal)
        } else if self.0.is_empty() {
            Some(std::cmp::Ordering::Greater)
        } else if other.0.is_empty() {
            Some(std::cmp::Ordering::Less)
        } else {
            self.0.peek().unwrap().partial_cmp(other.0.peek().unwrap())
        }
    }
}

impl<> Ord for HeapFragment {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        if self == other {
            std::cmp::Ordering::Equal
        } else if self.0.is_empty() {
            std::cmp::Ordering::Greater
        } else if other.0.is_empty() {
            std::cmp::Ordering::Less
        } else {
            self.0.peek().unwrap().cmp(other.0.peek().unwrap())
        }
    }
}

impl HeapHeap {
    fn max(&mut self) -> u32 {
        let mut h = self.heaps.pop().unwrap();
        let v = h.0.pop().unwrap();
        if !h.0.is_empty() {
            self.heaps.push(h);
        }
        v
    }
    fn merge(&mut self, mut other: BinaryHeap<HeapFragment>) {
        if other.is_empty() {
            return;
        }
        self.heaps.append(&mut other);
    }
    fn insert(&mut self, element: u32) {
        self.heaps.push(HeapFragment(BinaryHeap::from([element])));
    }
    fn new() -> Self {
        Self {
            heaps: BinaryHeap::new(),
        }
    }
}

fn main() {
    let mut input_file = BufReader::new(File::open("mergeheap.in").unwrap());
    let mut current_line = String::new();
    let _  = input_file.read_line(&mut current_line).unwrap();
    let (n, q) = current_line.trim().split_once(' ').unwrap();
    let n = n.parse::<u32>().unwrap();
    let q = q.parse::<u32>().unwrap();
    let mut hh: Vec<HeapHeap> = vec![HeapHeap::new(); n as usize];
    let mut out_file = BufWriter::new(File::create("mergeheap.out").unwrap());
    for _ in 0..q {
        current_line.clear();
        let _  = input_file.read_line(&mut current_line).unwrap();
        let (op, rest) = current_line.trim().split_once(' ').unwrap();
        match op {
            "1" => {
                let (m, x) = rest.split_once(' ').unwrap();
                let m = m.parse::<usize>().unwrap();
                let x = x.parse::<u32>().unwrap();
                hh[m - 1].insert(x);
            }
            "2" => {
                let m = rest.parse::<usize>().unwrap();                
                let val = hh[m - 1].max();
                let _ = writeln!(out_file, "{val}").unwrap();
            }
            "3" => {
                let (a, b) = rest.split_once(' ').unwrap();
                let a = a.parse::<usize>().unwrap();
                let b = b.parse::<usize>().unwrap();
                let b = hh[b - 1].heaps.drain().collect();
                hh[a - 1].merge(b);
            }
            _ => panic!("Unkown op {op}")
        }
    }
}