Cod sursa(job #2753085)

Utilizator PletoPletosu Cosmin-Andrei Pleto Data 20 mai 2021 23:27:46
Problema Paduri de multimi disjuncte Scor 0
Compilator rs Status done
Runda Arhiva educationala Marime 2.07 kb
use std::fs;
use std::io::{Write, BufRead};
use std::error::Error;

pub struct DisjointTree {
    parent: Vec<usize>
}

impl DisjointTree {
    pub fn new() -> Self {
        let mut vec = vec![0; 100001];
        for i in 0..100001 {
            vec[i] = i;
        }
        DisjointTree {
            parent: vec
        }
    }

    pub fn find(&mut self, x: usize) -> usize {
        if x == self.parent[x] {
            x
        } else {
            let parent = self.parent[x];
            self.parent[x] = self.find(parent);
            self.parent[x]
        }
    }

    pub fn unite(&mut self, x: usize, y: usize) {
        let upmost = self.find(x);
        self.parent[upmost] = y;
    }
}

fn main() -> Result<(), Box<dyn Error>> {
    let in_file = fs::File::open("disjoint.in")?;
    let out_file = fs::OpenOptions::new().append(true).open("disjoint.out")?;
    let mut in_buffer = std::io::BufReader::new(in_file);
    let mut out_buffer = std::io::BufWriter::new(out_file);

    let mut first_line = String::new();
    in_buffer.read_line(&mut first_line)?;

    let mut parts = first_line
        .split_whitespace()
        .map(|s| s.parse::<i32>());

    match (parts.next(), parts.next()) {
        (Some(Ok(_n)), Some(Ok(_m))) => {
            let mut disjoint = DisjointTree::new();

            in_buffer.lines()
                .map(|l| l.unwrap())
                .map(|s| {
                    s.split_whitespace()
                        .map(|i| i.parse::<usize>().unwrap())
                        .collect()
                })
                .for_each(|vec: Vec<usize>| {
                    let (cod, x, y) = (vec[0], vec[1], vec[2]);
                    if cod == 1 {
                        disjoint.unite(x, y);
                    } else {
                        match disjoint.find(x) == disjoint.find(y) {
                            true => out_buffer.write_all(b"DA\n"),
                            false => out_buffer.write_all(b"NU\n")
                        }.unwrap();
                    }
                })
        },
        _ => unreachable!()
    }

    Ok(())
}