Cod sursa(job #3228737)

Utilizator andreioneaAndrei Onea andreionea Data 10 mai 2024 20:13:13
Problema Heapuri cu reuniune Scor 100
Compilator rs Status done
Runda Arhiva educationala Marime 4.9 kb
use std::{
    fs::File,
    io::{BufRead, BufReader, BufWriter, Write},
};

#[derive(Debug, Clone)]
struct LinkedListNode<T> {
    element: T,
    next: Option<Box<LinkedListNode<T>>>,
}

impl<T> LinkedListNode<T> {
    fn new(element: T) -> Self {
        Self {
            element,
            next: None,
        }
    }
}

#[derive(Debug, Clone)]
struct LinkedList<T> {
    head: Option<Box<LinkedListNode<T>>>,
}

impl<T> LinkedList<T> {
    fn new() -> Self {
        Self {
            head: None,
        }
    }
}

impl<T> LinkedList<T> {
    fn drain(&mut self) -> Self {
        Self {
            head: self.head.take(),
        }
    }
    fn is_empty(&self) -> bool {
        self.head.is_none()
    }
    fn pop(&mut self) -> Option<Box<LinkedListNode<T>>> {
        if self.is_empty() {
            return None;
        }
        let next = self
            .head
            .as_mut()
            .unwrap()
            .next
            .take();
        let mut to_return = self.head.take();
        self.head = next;
        to_return.as_mut().unwrap().next = None;
        to_return
    }
    fn push(&mut self, element: T) {
        let mut new_node = Box::new(LinkedListNode::new(element));
        if self.is_empty() {
            self.head = Some(new_node);
        } else {
            new_node.as_mut().next = self.head.take();
            self.head = Some(new_node);
        }
    }
}

#[derive(Debug, Clone)]
struct PairingHeap {
    top: Option<u32>,
    sub_heaps: LinkedList<PairingHeap>,
}

impl PairingHeap {
    fn drain(&mut self) -> Self {
        Self {
            top: self.top.take(),
            sub_heaps: self.sub_heaps.drain(),
        }
    }
    fn is_empty(&self) -> bool {
        self.top.is_none() && self.sub_heaps.is_empty()
    }

    fn max(&mut self) -> Option<u32> {
        let val = self.top.take();
        let mut sub_heaps = self.sub_heaps.drain();
        while !sub_heaps.is_empty() {
            let mut child = sub_heaps.pop();
            if child.is_some() {
                self.meld(&mut (child.as_mut().unwrap().element));
            }
        }
        val
    }
    fn meld(&mut self, other: &mut PairingHeap) {
        if other.is_empty() {
            return;
        }
        if self.is_empty() {
            self.top = other.top;
            self.sub_heaps = other.sub_heaps.drain();
            return;
        }
        if self.top.unwrap() >= other.top.unwrap() {
            self.sub_heaps.push(other.drain());
        } else {
            let old_top = self.top.take();
            let old_sub_heaps = self.sub_heaps.drain();
            let old_self = PairingHeap {
                top: old_top,
                sub_heaps: old_sub_heaps,
            };
            self.top = other.top.take();
            self.sub_heaps = other.sub_heaps.drain();
            self.sub_heaps.push(old_self);
        }
    }
    fn insert(&mut self, element: u32) {
        let mut other = PairingHeap {
            top: Some(element),
            sub_heaps: LinkedList::new(),
        };
        self.meld(&mut other);
    }
}

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 ph: Vec<PairingHeap> = vec![
        PairingHeap {
            top: None,
            sub_heaps: LinkedList::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();
                ph[m - 1].insert(x);
            }
            "2" => {
                let m = rest.parse::<usize>().unwrap();
                let val = ph[m - 1].max().unwrap();
                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();
                // Can't do two mutable borrows at the same time from a vector :sadcat:
                if a < b {
                    let (aa, bb) = ph.split_at_mut(b - 1);
                    aa[a - 1].meld(&mut bb[0]);
                } else if a > b {
                    let (aa, bb) = ph.split_at_mut(a - 1);
                    bb[0].meld(&mut aa[b - 1]);
                } else {
                    // WAT
                }
            }
            _ => panic!("Unknown op {op}"),
        }
    }
}