Cod sursa(job #3364939)

Utilizator tryharderulbrebenel mihnea stefan tryharderul Data 14 septembrie 2026 09:29:02
Problema Paduri de multimi disjuncte Scor 100
Compilator rs Status done
Runda Arhiva educationala Marime 2.94 kb
use std::fs::File;
use std::io::{Read, Write};

struct DSU {
    // Using u32 instead of usize halves memory usage and improves CPU cache locality
    t: Vec<u32>,
    sz: Vec<u32>,
}

impl DSU {
    fn new(n: usize) -> DSU {
        let sz = vec![1; n + 1];
        // Efficient way to initialize the parent array
        let t: Vec<u32> = (0..=n as u32).collect();
        DSU { t, sz }
    }

    // Iterative find is generally faster than recursive and has zero risk of Stack Overflow
    fn find(&mut self, mut node: u32) -> u32 {
        let mut root = node;
        // 1. Find the root
        while root != self.t[root as usize] {
            root = self.t[root as usize];
        }

        // 2. Path compression (flatten the tree)
        while node != root {
            let nxt = self.t[node as usize];
            self.t[node as usize] = root;
            node = nxt;
        }

        root
    }

    fn unite(&mut self, mut x: u32, mut y: u32) {
        x = self.find(x);
        y = self.find(y);

        if x != y {
            // Union by size
            if self.sz[x as usize] < self.sz[y as usize] {
                // Swap without std::mem::swap to avoid overhead
                let temp = x;
                x = y;
                y = temp;
            }
            self.sz[x as usize] += self.sz[y as usize];
            self.t[y as usize] = x;
        }
    }
}

fn main() {
    // 1. FAST INPUT: Read all bytes to avoid String overhead
    let mut input_bytes = Vec::new();
    File::open("disjoint.in")
        .unwrap()
        .read_to_end(&mut input_bytes)
        .unwrap();

    let mut cursor = 0;

    let mut next_num = || -> u32 {
        while cursor < input_bytes.len() && input_bytes[cursor] <= b' ' {
            cursor += 1;
        }
        if cursor == input_bytes.len() {
            return 0;
        }
        let mut res = 0;
        while cursor < input_bytes.len() && input_bytes[cursor] > b' ' {
            res = res * 10 + (input_bytes[cursor] - b'0') as u32;
            cursor += 1;
        }
        res
    };

    let n = next_num() as usize;
    if n == 0 { return; }
    let m = next_num() as usize;

    let mut ds = DSU::new(n);

    // 2. FAST OUTPUT: Pre-allocate a byte vector for all answers.
    // At most `m` operations, each printing max 3 bytes ("DA\n" or "NU\n").
    let mut output_bytes = Vec::with_capacity(m * 3);

    for _ in 0..m {
        let op = next_num();
        let x = next_num();
        let y = next_num();

        if op == 1 {
            ds.unite(x, y);
        } else if op == 2 {
            if ds.find(x) == ds.find(y) {
                // extend_from_slice writes bytes directly to the vector
                output_bytes.extend_from_slice(b"DA\n");
            } else {
                output_bytes.extend_from_slice(b"NU\n");
            }
        }
    }

    // Dump all answers to the file at once
    File::create("disjoint.out")
        .unwrap()
        .write_all(&output_bytes)
        .unwrap();
}