ABC070 の D 問題を python で解いてみる

木構造の距離を求める問題演習として, AtCoder Beginner Contest 070 の D 問題の解説をしてみます.

グラフの中でも木構造というのはかなり単純な構造になっていまして, 無向グラフであって連結かつ閉路を含まないという性質から,
1. 頂点の数を $N$ とすると, 辺の数は $N-1$ となっている.
2. 2点間のパスは1通りしかない.
等の性質があります. 一般にグラフの最短距離を求めるアルゴリズムは再帰に工夫が必要になりますが, 木構造だと一度目標の頂点たどり着けばそれで良いというウレシイ状態になっています.

さてこの問題ですが, 木構造では2点間パスが1通りしかないという性質から, 頂点 $x_j$ から頂点 $K$ を経由しつつ頂点 $y_j$ に移動する場合の最短経路は,
$\text{"頂点 }x_j \text{ から頂点 }K\text{ までの距離"}+\text{"頂点 }K \text{ から頂点 }y_j\text{ までの距離"}$
によって求めることが出来ます.
ところが, 頂点の数と辺の数は $N\sim 10^5$ であり, 質問クエリも $Q\sim 10^5$ 個あるので質問に対して逐一計算するなると総計算量が $O(NQ)$ となり時間内に計算し終わることが出来ません. そこで, 頂点 $K$ からの各頂点までの距離を1度に計算してしまい, それをメモすることによって時間を省略してしまいます.

まずは入力編

N=int(input())
edges=[list(map(int,input().split())) for i in range(N)]
Q,K=map(int,input().split())
xy=[list(map(int,input().split())) for i in range(Q)]

メモを行うため, 前回のアルゴリズムを少し修正します.

def BFS(K,edges,N):
  roots=[ [] for i in range(N)]
  for a,b,c in edges:
    roots[a-1]+=[(b-1,c)]
    roots[b-1]+=[(a-1,c)]
  dist=[-1]*N
  stack=[]
  stack.append(K)
  dist[K]=0
  while stack:
    label=stack.pop(-1)
    for i,c in roots[label]:
      if dist[i]==-1:
        dist[i]=dist[label]+c
        stack+=[i]
  return dist

前回のアルゴリズムでは終点に到着した場合にそこまでの距離を返していましたが, 今回は未到達の点がなくなるまで計算を行い, リスト $\textrm{dist}$ そのものを返します.
これによって $\textrm{BFS}(v_0,\textrm{edges},N)$ の第 $i-1$ 成分は頂点 $v_0$ から頂点 $i$ までの距離を表しています.

あとは出力編

distance=BFS(K-1,edges,N)
for i in range(Q):
  print(distance[xy[i][0]-1]+distance[xy[i][1]-1])

と, 単なる代入によって質問クエリに答えていけばよいです. 以上をまとめて,

def BFS(K,edges,N):
  roots=[ [] for i in range(N)]
  for a,b,c in edges:
    roots[a-1]+=[(b-1,c)]
    roots[b-1]+=[(a-1,c)]
  dist=[-1]*N
  stack=[]
  stack.append(K)
  dist[K]=0
  while stack:
    label=stack.pop(-1)
    for i,c in roots[label]:
      if dist[i]==-1:
        dist[i]=dist[label]+c
        stack+=[i]
  return dist
 
N=int(input())
edges=[]
for i in range(N-1):
  edges+=[list(map(int,input().split()))]
Q,K=map(int,input().split())
xy=[list(map(int,input().split())) for i in range(Q)]
distance=BFS(K-1,edges,N)
for i in range(Q):
  print(distance[xy[i][0]-1]+distance[xy[i][1]-1])

を submit すれば AC となります.
Submission #1569148 - AtCoder Beginner Contest 070 | AtCoder

また, 今回は入力の個数も非常に多いため, input の前に次のようなおまじないをすると時間が短縮されます.

from sys import stdin
input = stdin.readline

Submission #1565687 - AtCoder Beginner Contest 070 | AtCoder