Cod sursa(job #3228801)

Utilizator andreioneaAndrei Onea andreionea Data 11 mai 2024 12:10:47
Problema Lowest Common Ancestor Scor 70
Compilator rs Status done
Runda Arhiva educationala Marime 4.63 kb
use std::{collections::HashMap, fs::File, io::{BufRead, BufReader, BufWriter, Write}};

struct DSU {
    parent: Vec<u32>,
    sizes: Vec<u32>,
}

impl DSU {
    fn new(set_size: usize) -> Self {
        let mut parent: Vec<u32> = Vec::with_capacity(set_size as usize);
        for i in 0..set_size {
            parent.push(i as u32);
        }
        Self {
            parent,
            sizes: vec![0; set_size]
        }
    }
    fn find_set(&mut self, u: u32) -> u32 {
        if self.parent[u as usize] == u {
            u
        } else {
            let result = self.find_set(self.parent[u as usize]);
            // Path compression.
            self.parent[u as usize] = result;
            result
        }
    }
    fn union(&mut self, a: u32, b: u32) {
        let a_set = self.find_set(a);
        let b_set = self.find_set(b);
        if a_set != b_set {
            if self.sizes[a_set as usize] > self.sizes[b_set as usize] {
                self.parent[b_set as usize] = a_set;
                self.sizes[a_set as usize] += self.sizes[b_set as usize];
            } else {
                self.parent[a_set as usize] = b_set;
                self.sizes[b_set as usize] += self.sizes[a_set as usize];
            }
        }
    }
}

struct Tree {
    parents: Vec<u32>,
    children: Vec<Vec<u32>>
}

impl Tree {
    fn new(n: usize, parents: Vec<u32>) -> Self {
        let mut children = vec![vec![]; n];
        for (child, parent) in parents.iter().enumerate() {
            if child == *parent as usize {
                continue;
            }
            children[*parent as usize].push(child as u32);
        }
        Self {
            parents,
            children
        }
    }
}

struct Query {
    other_node: u32,
    query_id: usize
}

struct TarjanLCAQuerySolver {
    queries_by_node: HashMap<u32, Vec<Query>>,
    answers: Vec<u32>,
    ancestors: Vec<u32>,
    visited: Vec<bool>,
    dsu: DSU,
    tree: Tree
}

impl TarjanLCAQuerySolver {
    fn new(tree: Tree, queries: Vec<(u32, u32)>) -> Self {
        let m = queries.len();
        let n = tree.parents.len();
        let mut queries_by_node = HashMap::with_capacity(n);
        for id in 0..n {
            queries_by_node.insert(id as u32, vec![]);
        }
        for (id, (u,v)) in queries.into_iter().enumerate() {
            queries_by_node.get_mut(&u).unwrap().push(Query{other_node: v, query_id: id});
            queries_by_node.get_mut(&v).unwrap().push(Query{other_node: u, query_id: id});
        }
        Self {
            queries_by_node,
            answers: vec![0; m],
            ancestors: vec![0; n],
            visited: vec![false; n],
            dsu: DSU::new(n),
            tree
        }
    }
    fn dfs(&mut self, node: u32) {
        self.visited[node as usize] = true;
        self.ancestors[node as usize] = node;
        let to_visit = self.tree.children[node as usize].clone();
        for child in to_visit {
            if !self.visited[child as usize] {
                self.dfs(child);
                self.dsu.union(child, node);
                self.ancestors[self.dsu.find_set(node) as usize] = node;
            }
        }
        if let Some(queries) = self.queries_by_node.get(&node) {
            for q in queries.iter() {
                if self.visited[q.other_node as usize] {
                    self.answers[q.query_id] = self.ancestors[self.dsu.find_set(q.other_node) as usize] + 1;
                }
            }
        }
    }
    fn solve(&mut self) {
        self.dfs(0);
    }
}

fn main() {
    let mut input_file = BufReader::new(File::open("lca.in").unwrap());
    let mut line = String::new();
    let _ = input_file.read_line(&mut line).unwrap();
    let (n, m) = line.trim().split_once(' ').unwrap();
    let n = n.parse::<usize>().unwrap();
    let m = m.parse::<usize>().unwrap();
    let mut parents = Vec::with_capacity(n);
    parents.push(0);
    line.clear();
    let _ = input_file.read_line(&mut line).unwrap();
    for parent in line.trim().split(' ') {
        let parent = parent.parse::<u32>().unwrap() - 1;
        parents.push(parent);
    }
    let tree = Tree::new(n, parents);
    let mut queries = Vec::with_capacity(m);
    for _ in 0..m {
        line.clear();
        let _ = input_file.read_line(&mut line).unwrap();
        let (u,v) = line.trim().split_once(' ').unwrap();
        let u = u.parse::<u32>().unwrap() - 1;
        let v = v.parse::<u32>().unwrap() - 1;
        queries.push((u,v));

    }
    let mut solver = TarjanLCAQuerySolver::new(tree, queries);
    solver.solve();
    let mut out_file = BufWriter::new(File::create("lca.out").unwrap());
    for ans in solver.answers {
        writeln!(out_file, "{ans}").unwrap();
    }
}