ABC 049, abc done, d is too slow and incorrect, need to learn more about union find

atcoder/beginner_contests/abc049/src/bin/d.rs

@@ -0,0 +1,76 @@
+use proconio::input;
+use std::collections::{HashSet, VecDeque};
+
+fn dfs(start: usize, vertices: &[Vec<usize>]) -> HashSet<usize> {
+    let mut result = HashSet::new();
+    let mut visited = HashSet::new();
+    let mut queue = VecDeque::new();
+    queue.push_back(start);
+
+    while let Some(current_vertex) = queue.pop_front() {
+        result.insert(current_vertex);
+
+        for neighbour in &vertices[current_vertex][..] {
+            if visited.insert(*neighbour) {
+                queue.push_front(*neighbour);
+            }
+        }
+    }
+
+    result
+}
+
+fn main(){
+    input!{
+        n: usize,
+        k: usize,
+        l: usize,
+        pq: [(usize, usize); k],
+        rs: [(usize, usize); l],
+    };
+
+    let mut roads = vec![Vec::new(); n];
+    for (p, q) in pq {
+        roads[p-1].push(q-1);
+    }
+
+    let mut rails = vec![Vec::new(); n];
+    for (r, s) in rs {
+        rails[r-1].push(s-1);
+    }
+
+    let mut city_found_road = Vec::new();
+    let mut city_to_road_subsets = vec![-1; n];
+    for city in 0..n {
+        if city_to_road_subsets[city] == -1 {
+            let connected_cities = dfs(city, &roads);
+            let l = city_found_road.len();
+            for connected_city in &connected_cities {
+                city_to_road_subsets[*connected_city] = l as i32;
+            }
+            city_found_road.push(connected_cities);
+        }
+    }
+
+    let mut city_found_rails = Vec::new();
+    let mut city_to_rail_subsets = vec![-1; n];
+    for city in 0..n {
+        if city_to_rail_subsets[city] == -1 {
+            let connected_cities = dfs(city, &rails);
+            let l = city_found_rails.len();
+            for connected_city in &connected_cities {
+                city_to_rail_subsets[*connected_city] = l as i32;
+            }
+            city_found_rails.push(connected_cities);
+        }
+    }
+
+    for city in 0..n {
+        let rail_set = &city_found_rails[city_to_rail_subsets[city] as usize];
+        let road_set = &city_found_road[city_to_road_subsets[city] as usize];
+
+        let result : HashSet<_> = rail_set.intersection(road_set).collect();
+
+        print!("{} ", result.len());
+    }
+}