ABC214 を Rust で解く (D問題まで)

# はじめに

[AtCoder Beginner Contest 214](https://atcoder.jp/contests/abc214) の D問題までを Rust で解きました！
詳しめに解説をしていこうと思います。
役にたちましたら LGTM いただけると嬉しいです。
AC 確認はしてますが、もっとスマートな書き方がありましたら教えてください。


# なぜやるのか

コンテストから時間をあけて解きなおすことで、記憶へ定着すればいいなと思っています。
また、会社で Rust を使うのですが私はまだ慣れておらず、その練習のためにやっています。
(コンテスト本番には C++ を使って参加している水色コーダーです。)


# 目次

[A - New Generation ABC](#a---new-generation-abc) (<font color="#808080">Difficulty 4</font>)
[B - How many?](#b---how-many) (<font color="#808080">Difficulty 53</font>)
[C - Distribution](#c---distribution) (<font color="#808080">Difficulty 319</font>)
[D - Sum of Maximum Weights](#d---sum-of-maximum-weights) (<font color="#00C0C0">Difficulty 1341</font>)
[UnionFind木](#unionfind木)

# 前提

この記事にあるコードはこちらを import しています。

```rust
#[allow(unused_imports)]
use proconio::{input, fastout, marker::Chars};
#[allow(unused_imports)]
use std::collections::{HashSet, HashMap, BTreeSet, VecDeque};
#[allow(unused_imports)]
use std::cmp::{max, min};
#[allow(unused_imports)]
use itertools::Itertools;
```


# A - New Generation ABC

問題文に条件式があるので、そのままプログラムにしましょう。
境界値に注意です。

```rust
fn main() {
  input! {
    n: usize
  }
  let ans = if n <= 125 {
    4
  } else if n <= 211 {
    6
  } else {
    8
  };
  println!("{}", ans);
}
```


# B - How many? 

問題が [ARC113-A](https://atcoder.jp/contests/arc113/tasks/arc113_a) に似ているので少し身構えました。
全探索をすると 3重ループになって $O(S^3)$ ですが、制約の値が小さいので解けますね。
[ARC113-A](https://atcoder.jp/contests/arc113/tasks/arc113_a) ではこれに似た問題をもっと高速に解くので、見てみるのも面白いと思います。

```rust
fn main() {
  input! {
    s: usize,
    t: usize
  }
  let mut ans = 0;
  for a in 0..=s {
    for b in 0..=s {
      for c in 0..=s {
        if a + b + c > s { break; }
        if a * b * c <= t { ans += 1; } 
      }
    }
  }
  println!("{}", ans);
}
```


# C - Distribution 

少し複雑ですね。
まずは問題文にある挙動を考えます。つまり、宝石をもらった時間が早い順に処理することを考えます。
すると、最大でも $2 * N$ 回の処理をすればいいことがわかりますので、この方針で解きます。
最小値にだけ用があるので、こういうときは `Heap` をつかいましょう。 Rust では `BinaryHeap` という名前です。

まず、時刻 $T_i$ に確定で宝石が渡されるので、これを `Heap` にいれます。答えを記録しておくためにインデックスが必要なので、時刻とインデックスのペアにしています。

あとは、時間の小さい順に `Heap` から取り出して答えを更新します。
更新した後に $S_i$ を足して `Heap` に入れなおすことと、この時のインデックスは `(idx + 1) % n` であることに注意です。

```rust
use std::cmp::Reverse;
use std::collections::BinaryHeap;

fn main() {
  input! {
    n: usize,
    s: [usize; n],
    t: [usize; n]
  }

  let mut ans = vec![1001001001; n];
  let mut que = BinaryHeap::new();
  let t = t.into_iter().enumerate().collect_vec();
  for (idx, val) in t.into_iter() { //Ti を heap にいれる
    que.push(Reverse((val, idx)));
  }

  while let Some(Reverse((val, idx))) = que.pop() {
    if ans[idx] > val {
      ans[idx] = val;
      que.push(Reverse((val + s[idx], (idx + 1) % n))); // Si の時間を足して、 idx + 1 として heap にいれる
    }
  }

  for i in 0..n {
    println!("{}", ans[i]);
  }
}
```


# D - Sum of Maximum Weights 

グラフの問題ですね。難しいです。

$\sum_{i=1}^{N-1} \sum_{j=i+1}^N$ の時点で $O(N^2)$ ですね。 $2 \le N \le 10^5$ なので、このままでは間に合いません。
$\sum$ があるときは主客転倒を考えます。つまり、 $f(u_i, v_j) = w_i$ となる $(i, j)$ の組み合わせがいくつあるかを考えます。
この個数を $c_i$ とすると、 $\sum_{i=1}^N w_i * c_i$ が答えになります。これなら、 $c_i$ を 遅くとも $O(logN)$ 程度で見つけられれば時間内に解けそうです。

さて、 $c_i$ を考えましょう。$w_i$ の重みをもつ辺に注目します。この辺がつないでいる2頂点を $v_1,\ v_2$ とすると、
$c_i =$ [ $v_1$ 側の部分木で、重さが $w_i$ 以下の枝でつながっているものの個数 ] * [ $v_2$ 側の部分木で、重さが $w_i$ 以下の枝でつながっているものの個数 ] 
だとわかります。私は紙で試していると浮かびました (再現性がなくてすみません)。

「$w_i$ 以下」が少し考えにくいので、 $w_i$ は全ての枝の中で最大値だとしましょう。そうすると単純になって、次のように言い換えられます。
[ 重さが $w_i$ の辺でグラフを2つに分けた時、それぞれのグラフの連結成分の個数の積 ]
こうすると嬉しいのは、実際にグラフをカットしていくことで、 $w_i$ が大きい順に処理できることです。
さて、こうするとグラフの連結成分の個数を求めれば答えがわかりますね。ですが問題があって、連結成分の個数を求めるには $O(N)$ かかります。このとき、全体の計算量が $O(N^2)$ になるので時間内に間に合いません。

この部分、もっと速くできないでしょうか。できます。

上で考えたアルゴリズムの挙動を考えると、グラフがどんどん分解されて、最後は辺がなくなって頂点だけになります。この挙動、何かに似ていませんか。そうです。 UnionFind木です。[ABC120-D](https://atcoder.jp/contests/abc120/tasks/abc120_d) に似ているので、解いたことある方はひらめきやすいかと思います。
UnionFind木は連結していくのですが、今回は分解していくので、挙動が逆向きです。なので、今回の問題も分解ではなく連結していく方向で解けないかを考えます。すると、先に考えたアルゴリズムを以下のように変更すると解けそうです。

・$w_i$ を最大値からみる -> $w_i$ を最小値から見る
・グラフをカットする -> グラフを連結する
・連結成分の個数を $O(N)$ で計算 -> 連結成分の個数を $O(\alpha(N))$ で計算

こうすると、計算量が $O(N * \alpha(N))$ になりました。
($\alpha(N)$ はアッカーマン関数の逆関数で、とても小さな値になるので定数倍とみなせます)
ソートがあるので全体では $O( N log N )$ になって、これは時間内に解けます。

実装です。


```rust
fn main() {
  input! {
    n: usize,
    mut uvw: [(usize, usize, usize); n - 1]
  }

  uvw.sort_by_key(|v| v.2); // 重さ w で昇順ソート
  let mut uf = UnionFind::new(n);
  let mut ans = 0;
  for (u, v, w) in uvw.into_iter() {
    let u = u - 1;
    let v = v - 1;
    ans += w * uf.size(u) * uf.size(v);
    uf.unite(u, v);
  }
  println!("{}", ans);
}
```

使用した UnionFind木は以下に載せます。

# UnionFind木

```rust
struct UnionFind {
  p: Vec<usize>, // 親
  rank: Vec<usize>	// サイズ・各集合の根のみ有効
}

#[allow(dead_code)]
impl UnionFind {
  fn new(size: usize) -> Self {
    UnionFind{
      p: (0..size).collect(),
      rank: vec![1; size]
    }
  }

  fn root(&mut self, x: usize) -> usize {
    if self.p[x] == x {
      x
    }
    else {
      self.p[x] = self.root(self.p[x]);
      self.p[x]
    }
  }

  fn unite(&mut self, x: usize, y: usize)  {
    let mut x = self.root(x);
    let mut y = self.root(y);
    if x == y { return /*false*/ }

    if self.rank[x] > self.rank[y] { std::mem::swap(&mut x, &mut y); }
    self.rank[y] += self.rank[x];
    self.p[x] = y;
    // true
  }

  fn root_num(self) -> usize {
    self.p.iter().enumerate().filter(|(idx, val)| idx == *val).count()
  }

  fn groups(self) -> std::collections::HashMap<usize, Vec<usize>> {
    let mut mp = std::collections::HashMap::new();
    for i in 0..self.p.len() {
      (*mp.entry(i).or_insert(vec![])).push(i);
    }
    mp
  }

  fn size(&mut self, x: usize) -> usize {
    let root_x = self.root(x);
    self.rank[root_x]
  }

  fn same(&mut self, x: usize, y: usize) -> bool {
    self.root(x) == self.root(y)
  }
}
```

----

最後まで読んでいただきありがとうございます！