Add an example for ABC157-D

qryxip · qryxip · commit 2445d012cff4 · 2020-05-19T00:20:09.000+09:00
diff --git a/examples/abc157-d.rs b/examples/abc157-d.rs
@@ -0,0 +1,59 @@
+// https://atcoder.jp/contests/abc157/tasks/abc157_d
+//
+// 以下のクレートを使用。
+//
+// - `itertools`
+// - `petgraph`
+// - `proconio`
+
+use itertools::Itertools as _;
+use petgraph::graphmap::UnGraphMap;
+use petgraph::unionfind::UnionFind;
+use proconio::input;
+use proconio::marker::Usize1;
+
+fn main() {
+    // `proconio::input!`。
+    //
+    // https://docs.rs/proconio/0.3.6/proconio/macro.input.html
+    input! {
+        n: usize,
+        m: usize,
+        k: usize,
+        abs: [(Usize1, Usize1); m],
+        cds: [(Usize1, Usize1); k],
+    }
+
+    // `petgraph::unionfind::UnionFind<_>`を使用。
+    // `UnionFind`は各集合のサイズは覚えてないので`vec![_; n]`で管理する。
+    //
+    // https://docs.rs/petgraph/0.5/petgraph/unionfind/struct.UnionFind.html
+    let mut ff = UnionFind::new(n);
+    let mut ff_sizes = vec![1; n];
+    for &(a, b) in &abs {
+        let size_sum = ff_sizes[ff.find(a)] + ff_sizes[ff.find(b)];
+        if ff.union(a, b) {
+            ff_sizes[ff.find(a)] = size_sum;
+        }
+    }
+
+    // `UnionFind`の`rank`はprivateな為面倒。ここではこのような方法で数える。
+    // この際`HashSet`のかわりに`petgraph::graphmap::GraphMap<..>`を使う。
+    //
+    // https://docs.rs/petgraph/0.5/petgraph/graphmap/struct.GraphMap.html
+    let mut ans = (0..n).map(|i| ff_sizes[ff.find(i)] - 1).collect::<Vec<_>>();
+    let mut already_decr = UnGraphMap::with_capacity(n, abs.len());
+    for (a, b) in abs {
+        ans[a] -= 1;
+        ans[b] -= 1;
+        already_decr.add_edge(a, b, ());
+    }
+    for (c, d) in cds {
+        if ff.equiv(c, d) && !already_decr.contains_edge(c, d) {
+            ans[c] -= 1;
+            ans[d] -= 1;
+        }
+    }
+
+    println!("{}", ans.iter().format(" "));
+}
diff --git a/test-examples.toml b/test-examples.toml
@@ -184,6 +184,13 @@ url = "https://atcoder.jp/contests/abc156/tasks/abc156_d"
 matching = "Words"
 meta = { using = ["num", "proconio"] }
 
+[examples.abc157-d]
+type = "Normal"
+name = "ABC157: D - Friend Suggestions"
+url = "https://atcoder.jp/contests/abc157/tasks/abc157_d"
+matching = "Words"
+meta = { using = ["itertools", "petgraph", "proconio"] }
+
 [examples.agc023-a]
 type = "Normal"
 name = "AGC023: A - Zero-Sum Ranges"
