Pagini recente » Cod sursa (job #3036672) | Cod sursa (job #66335) | Cod sursa (job #1959416) | Cod sursa (job #1206062) | Cod sursa (job #3228801)
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();
}
}