ABC428 E - Farthest Vertex
https://atcoder.jp/contests/abc428/tasks/abc428_e

コンテスト中の動き

コンテスト中に少しやろうとして「木の直径」というキーワードだけが頭に思い浮かび、求め方までは知っているのですが一応確認するためにWebで検索しました。しかしそれ以上は何も思い浮かばなかったのでさっさと捨てました。……今思えば検索して出てくるサイトに

※ 「どの頂点から始めても、最も遠くにある頂点 x は直径の端点になる」という事実を利用しています。

割とヒントになることが書いてありましたね。コンテスト中の検索で何かに気づいたことってほぼ無いんですが、たぶん焦っているせいで僕には何も見えていないのでしょう。

方針

公式解説を読んでも「木に含まれる任意の頂点 u について、s と t の少なくとも一方が u から最も遠い頂点である。」の証明がよくわからなかったので公式解説放送を見ました。これによってようやく腑に落ちました。自分なりにまとめます。

木の直径と木の中心

木の中で作れる最長パスの長さを木の直径、その真ん中にある頂点や辺のことを木の中心と呼ぶようです。直径の長さが偶数なら中心は頂点、奇数なら中心は辺になりますね。解説放送ではこの2つの場合に分けて説明をされていました。

これを利用して公式解説にあるように、各頂点の最遠点が直径の端点のいずれかになるという話を考えます。

木の中心が辺になる場合

これは辺が木の中心になる場合です。この木はこの辺の両端点をそれぞれ根とする2つの部分木に分けることができます。図でいうとグループAに所属する頂点は、グループBの右端が最遠点になりそうです。グループBに所属する頂点はグループAの左端が最遠点になりそうです。直感的にはわかりますがもう少しだけ厳密に考えます。

各グループの中でもっと遠い点を用意できないか？ということです。この木の直径を $d$ とすると、グループA, グループBは図の中央の2点を根とする深さ $\frac{d-1}{2}$ の木になります。

図で赤く塗った頂点からの距離を考えてみます。グループBにおけるこの頂点の深さを $e$ とすると、グループAの左端の頂点までの距離は $\frac{d-1}{2} + 1 + e$ となります。また、グループBの中で最も長いパスを作ろうとすると、この頂点からグループBの根まで行って、そこからグループBの右端の頂点まで行くことになります。すなわち $e + \frac{d-1}{2}$ です。当然ながら

$$ \frac{d-1}{2} + 1 + e > e + \frac{d-1}{2} $$

なので、グループAの左端の方が遠くなります。

木の中心が頂点になる場合

木の中心が頂点になる場合を考えます。こんな感じですね。これもさっきと同じで、今度はそれぞれのグループの木の深さが $\frac{d}{2} - 1 $ になります。

所属していないグループの一番深い頂点までの距離は $\frac{d}{2} + 1 + e$、所属するグループ内での最遠点までの距離は $\frac{d}{2} - 1 + e$ になるので、やはり所属グループ外の最遠点がその頂点から見たときの最遠点になります。

$$ \frac{d}{2} + 1 + e > \frac{d}{2} - 1 + e$$

というわけで、いずれにしてもある頂点から見たときの最遠点は木の中心からの最遠点のどこかになります。

ここに書かれている木の直径を求めるアルゴリズム

任意の頂点 s を選ぶ

s から DFS/BFS によって、最も遠くにある頂点 u を探索する

u から DFS/BFS によって、最も遠くにある頂点 v を探索する

u と v を結ぶパスが直径となる

はこういう理屈でできているんですね。木の中心を用いた丁寧な解説放送のおかげで納得することができました。とてもありがたいです。

この性質を利用して、各頂点から他の全ての頂点への距離を計算して最遠点を求めるのではなく、最遠点の候補となる2点だけに着目してそこから出発して探索することにより各頂点への距離を求めます。$N^2$ 回の探索が $2 N$ 回の探索で済むイメージですね。

実装

この問題では最遠点が複数ある場合、最も頂点番号の大きなものを答えなくてはいけません。というわけで先に挙げたアルゴリズムを少し改造して、こうします。

任意の頂点 s を選ぶ
s から DFS/BFS によって、最も遠くにある頂点の中でも番号が最大である u を探索する
u から DFS/BFS によって、最も遠くにある頂点の中でも番号が最大である v を探索する
u と v を結ぶパスが直径となる

DFS/BFS によってある頂点から他の全ての頂点への距離が求まりますから、それらの距離をいったん配列に収めた後で条件を満たす u, v を求めればいいです。

さらに v からの DFS/BFS も行うことにより「u から各頂点への距離」「v から各頂点への距離」が出揃うので、後はこれらを比較するだけです。

というわけでこんな風に実装しています。

1 から BFS によって各頂点への距離を求め、最も遠くにある頂点の中でも番号が最大である s を求める
s から BFS によって各頂点への距離を求め、最も遠くにある頂点の中でも番号が最大である t を求める
t から BFS によって各頂点への距離を求める
1. と 3. で求めた距離を比較し、各頂点ごとにどちらが遠くなるのかを求める

from collections import deque
N = int(input())
G = [[] for _ in range(N+1)]
for _ in range(N-1):
  A, B = map(int, input().split())
  G[A].append(B)
  G[B].append(A)
# 木の直径の性質を利用して両端点を求める。
# 1 回目の BFS
que = deque([1]) 
dist1 = [-1 for _ in range(N+1)] 
dist1[1] = 0
while (que): 
  v = que.popleft() 
  for v2 in G[v]: 
    if dist1[v2] != -1: continue
    dist1[v2] = dist1[v] + 1
    que.append(v2)
# 1 からの最遠点 s を求める。同じ距離なら番号の大きい方を採用する。（vを昇順に見るので後の方が大きくなる）
max_dist, s = 0, 0 
for v in range(1, N+1): 
  if dist1[v] >= max_dist: 
    max_dist, s = dist1[v], v 
# 2 回目の BFS
que = deque([s]) 
dists = [-1 for _ in range(N+1)] 
dists[s] = 0
while (que): 
  v = que.popleft() 
  for v2 in G[v]: 
    if dists[v2] != -1: continue
    dists[v2] = dists[v] + 1
    que.append(v2)
# s からの最遠点 t を求める。同じ距離なら番号の大きい方を採用する。（vを昇順に見るので後の方が大きくなる）
max_dist, t = 0, 0 
for v in range(1, N+1): 
  if dists[v] >= max_dist: 
    max_dist, t = dists[v], v 
# 3 回目の BFSにより t からの各点への距離を求める。
que = deque([t]) 
distt = [-1 for _ in range(N+1)] 
distt[t] = 0
while (que): 
  v = que.popleft() 
  for v2 in G[v]: 
    if distt[v2] != -1: continue
    distt[v2] = distt[v] + 1
    que.append(v2)

# 実際に両端点 s, t までの距離を比べる。
for v in range(1, N+1):
  if dists[v] < distt[v]:
    print(t)
  elif dists[v] > distt[v]:
    print(s)
  else:
    print(max(s, t))