Official

E - MST + 1 Editorial by Nyaan


以下では \(G\) の辺集合を \(E\) と定義しています。また、頂点 \(u\) と頂点 \(v\) を結ぶ重み \(w\) の無向辺を \((u,v,w)\) と表しています。

まず、この問題をクエリごとに解く方法を説明します。 \(G_i\) の最小全域木 \(T_i\) はクラスカル法により以下の手順で \(O(M \log M)\) で求めることができます。

まず、\(N\) 頂点の UnionFind を用意する。
\(E \cup \lbrace e_i \rbrace\) を辺の重みでソートした後に、辺を重み昇順に見ていく。今見ている辺を \((u,v,w)\) として、

  • 頂点 \(u\) と頂点 \(v\) が UnionFind 上で連結ならば何もしない。
  • 頂点 \(u\) と頂点 \(v\) が UnionFind 上で非連結ならば、この辺を使用することにして UnionFind 上で \(u\)\(v\) を連結にする。

という操作をグラフが連結になるまで繰り返す。

上記のアルゴリズムによって \(T_i\) を得ることができるので、すべてのクエリに対して \(T_i\) を求めることによりこの問題の正答を得ることができますが、計算量が \(O(Q \times M \log M)\) かかり効率的とは言えません。

上の操作をよく見ると、クエリに答える ( = \(T_i\)\(e_i\) を含むかどうかを調べる) だけならば、\(e_i\) が出てきた時点でアルゴリズムを途中で打ち切ってもよいことがわかります。具体的には以下のような手順となります。(上のアルゴリズムと変化した部分は 太字 にしています。)

まず、N 頂点の UnionFind を用意する。
\(E \cup \lbrace e_i \rbrace\) を辺の重みでソートした後に、辺を重み昇順に見ていく。今見ている辺を \((u,v,w)\) として、

  1. 今見ている辺が \(e_i\) ならば
    • 頂点 \(u\) と頂点 \(v\) が非連結ならば Yes、連結ならば No を返し、操作を終了する。
  2. 今見ている辺が \(G\) の辺ならば
    • 頂点 \(u\) と頂点 \(v\) が UnionFind 上で連結ならば何もしない。
    • 頂点 \(u\) と頂点 \(v\) が UnionFind 上で非連結ならば、この辺を使用することにして UnionFind 上で \(u\)\(v\) を連結にする。

工夫によりアルゴリズムを少し改善できましたが、計算量はまだ \(O(Q \times M \log M)\) のままです。

上の操作をよく見ると、データ構造である UnionFind に影響を与えるのは \(G\) の辺のみで、\(e_i\) は影響を与えないことがわかります。
よって上記のアルゴリズムは、クエリ並列化、すなわちすべてのクエリ計算を並列に処理することが可能です。具体的には以下のアルゴリズムを実行しても\( O(N \log N)\) のアルゴリズムを \(Q\) 回実行するのと同じ結果になることが確かめられます。

まず、N 頂点の UnionFind を用意する。
\(E\) と全てのクエリで登場する辺、すなわち \(E \cup \lbrace e_1 , e_2 , \dots , e_Q \rbrace\) を辺の重みでソートした後に、辺を重み昇順に見ていく。今見ている辺を \((u,v,w)\) として、

  1. 今見ている辺が \(e_i\) \({(1 \leq i \leq Q)}\) ならば
    • クエリ \(i\) の答えを頂点 \(u\) と頂点 \(v\) が非連結ならば Yes、連結ならば No とする。
  2. 今見ている辺が \(G\) の辺ならば
    • 頂点 \(u\) と頂点 \(v\) が UnionFind 上で連結ならば何もしない。
    • 頂点 \(u\) と頂点 \(v\) が UnionFind 上で非連結ならば、この辺を使用することにして UnionFind 上で \(u\)\(v\) を連結にする。

上記のアルゴリズムは \(\mathrm{O}((M+Q) \log (M+Q))\) で動作して、この問題を解くのに十分な速度で動作します。
このように全てのクエリの答えを並列に求めるテクニックは \(500\) 点以上の問題でよく出てくるテクニックなので、クエリを与えられる問題が出たときは頭の片隅に入れておきましょう。

Python による解法は次の通りです。AC するには処理系に PyPy3 を使用する必要があります。また、ソースコードの71 行目までは ac-library-python を引用しています。

import typing

# reference
# not522, ac-library-python (Python port of AtCoder Library)
# https://github.com/not522/ac-library-python/blob/master/atcoder/dsu.py

class DSU:
  '''
  Implement (union by size) + (path halving)

  Reference:
  Zvi Galil and Giuseppe F. Italiano,
  Data structures and algorithms for disjoint set union problems
  '''

  def __init__(self, n: int = 0) -> None:
    self._n = n
    self.parent_or_size = [-1] * n

  def merge(self, a: int, b: int) -> int:
    assert 0 <= a < self._n
    assert 0 <= b < self._n

    x = self.leader(a)
    y = self.leader(b)

    if x == y:
      return x

    if -self.parent_or_size[x] < -self.parent_or_size[y]:
      x, y = y, x

    self.parent_or_size[x] += self.parent_or_size[y]
    self.parent_or_size[y] = x

    return x

  def same(self, a: int, b: int) -> bool:
    assert 0 <= a < self._n
    assert 0 <= b < self._n

    return self.leader(a) == self.leader(b)

  def leader(self, a: int) -> int:
    assert 0 <= a < self._n

    parent = self.parent_or_size[a]
    while parent >= 0:
      if self.parent_or_size[parent] < 0:
        return parent
      self.parent_or_size[a], a, parent = (
          self.parent_or_size[parent],
          self.parent_or_size[parent],
          self.parent_or_size[self.parent_or_size[parent]]
      )

    return a

  def size(self, a: int) -> int:
    assert 0 <= a < self._n

    return -self.parent_or_size[self.leader(a)]

  def groups(self) -> typing.List[typing.List[int]]:
    leader_buf = [self.leader(i) for i in range(self._n)]

    result: typing.List[typing.List[int]] = [[] for _ in range(self._n)]
    for i in range(self._n):
      result[leader_buf[i]].append(i)

    return list(filter(lambda r: r, result))


import sys

it = map(int, sys.stdin.buffer.read().split())
N, M, Q = next(it), next(it), next(it)
es = []
for u, v, w in zip(it, it, it):
  es.append((w, u, v, len(es)))

es.sort()
uf = DSU(N + 1)
answer = [None] * Q

for e in es:
  w, u, v, q = e
  if q >= M:
    if uf.same(u, v):
      answer[q - M] = "No"
    else:
      answer[q - M] = "Yes"
  else:
    uf.merge(u, v)

print(*answer, sep="\n")

発展問題:クエリが独立でない場合

この問題を解いた人の中には、クエリが後のクエリにも影響する場合にどうなるのかが気になった人がいるかもしれません。すなわちクエリが以下のような場合の問題です。

  • クエリ \(i\) : \(G\) に辺 \((u,v,w)\) を追加して、\((u,v,w)\) が最小全域木に使用されるかを YesNo で出力する。追加した辺は次のクエリ以降でも残っている。また、すべてのクエリでの辺の重みは相異なる。

発展問題は Union Find などの初等的なデータ構造だけで解くのは困難なのではないかと予想しています。 (高度なデータ構造なしでも分割統治で \(O(N \log^2 N)\) で解くことができるらしいですが自分はよくわかっていません。)

発展問題は Link/Cut Tree と呼ばれる高度なデータ構造を利用すれば解くことができます。 Link/Cut Tree とは、次の操作を償却計算量 \(O(\log N)\) で行うことができるデータ構造です。

  • 頂点 \(u\) と頂点 \(v\) が非連結であるとき、\(u-v\) 間に辺を追加する。(Link)
  • 頂点 \(u\) と頂点 \(v\) の間に辺があるとき、\(u-v\) 間の辺を削除する。(Cut)
  • 頂点 \(u\) と頂点 \(v\) が連結であるとき、頂点 \(u\), \(v\) 間のパス上の頂点が保持している値の最大値を取得 (「最大値」の部分は可換モノイドならばなんでもよい)

これらの操作が高速にできればクエリを償却 \(\mathrm{O}(\log N)\) で処理できるのはおそらく比較的容易に確認できると思います。(ちなみに、元の問題も Link/Cut Tree で \(\mathrm{O}(M \log M + Q \log N)\) で解くことができます。)

AtCoder で Link/Cut Tree が想定解法になることは 100 % ないと予想されるので AtCoder を楽しむ分には不要ですが、有志コンテストのクエリ問題では楽な非想定解法として登場する印象が強いので、コンテストで出題する側に回る人ならば詳しい話を知っておくといいことがあるかもしれません。

posted:
last update: