https://atcoder.jp/contests/abc304/tasks/abc304_e
UnionFindを作って、K個の条件はrootに直してペアをHashSetにして、クエリもrootに直してHashSetにあるかどうかを調べます。
RustではUnionFindをコーディングするのが大変ですね。parentも整数してやれば易しくなります。あと、手元では通ったusize::MAXがAtCoderでは見つかりませんでした。
// Good Graph #![allow(non_snake_case)] use std::cmp::max; use std::collections::HashSet; //////////////////// library //////////////////// fn read<T: std::str::FromStr>() -> T { let mut line = String::new(); std::io::stdin().read_line(&mut line).ok(); line.trim().parse().ok().unwrap() } fn read_vec<T: std::str::FromStr>() -> Vec<T> { read::<String>().split_whitespace() .map(|e| e.parse().ok().unwrap()).collect() } fn YesNo(b: bool) -> String { return if b { "Yes".to_string() } else { "No".to_string() } } //////////////////// Node //////////////////// type ID = usize; type Edge = (ID, ID); struct Node { id: ID, parent: ID, height: u32 } impl Node { const MAX: ID = 10000000; fn new(id: ID) -> Node { Node { id, parent: Node::MAX, height: 1 } } } //////////////////// UnionFind //////////////////// type UnionFind = Vec<Node>; fn make_tree(N: usize) -> UnionFind { (0..N).map(|i| Node::new(i)).collect::<UnionFind>() } fn find_root(tree: &UnionFind, v: ID) -> ID { let node = &tree[v]; if node.parent == Node::MAX { v } else { find_root(tree, node.parent) } } fn join(tree: &mut UnionFind, u: ID, v: ID) { let id1 = find_root(tree, u); let id2 = find_root(tree, v); if id1 == id2 { ; } else if tree[id1].height > tree[id2].height { tree[id2].parent = tree[id1].id; } else { tree[id1].parent = tree[id2].id; tree[id2].height = max(tree[id2].height, tree[id1].height + 1); } } //////////////////// process //////////////////// fn read_edge() -> Edge { let v: Vec<ID> = read_vec(); (v[0]-1, v[1]-1) } fn read_input() -> (usize, Vec<Edge>, Vec<(usize, usize)>, u32) { let v: Vec<usize> = read_vec(); let N = v[0]; let M = v[1]; let edges: Vec<Edge> = (0..M).map(|_| read_edge()).collect(); let K: u32 = read(); let conditions: Vec<(ID, ID)> = (0..K).map(|_| read_vec()). map(|v: Vec<ID>| (v[0]-1, v[1]-1)). collect(); let Q = read(); (N, edges, conditions, Q) } fn read_query() -> (ID, ID) { let v: Vec<ID> = read_vec(); (v[0]-1, v[1]-1) } fn f(N: usize, edges: Vec<Edge>, conditions: Vec<(usize, usize)>, Q: u32) { let mut tree = make_tree(N); for (u, v) in edges.into_iter() { join(&mut tree, u, v) } let pairs: HashSet<(ID, ID)> = conditions.into_iter(). map(|(u, v)| (find_root(&tree, u), find_root(&tree, v))). collect(); for _ in 0..Q { let (u, v) = read_query(); let r1 = find_root(&tree, u); let r2 = find_root(&tree, v); let b = pairs.contains(&(r1, r2)) || pairs.contains(&(r2, r1)); println!("{}", YesNo(!b)) } } //////////////////// main //////////////////// fn main() { let (N, edges, conditions, Q) = read_input(); f(N, edges, conditions, Q) }