AtCoder BC070D: Transit Tree Path

問題

https://beta.atcoder.jp/contests/abc070/tasks/abc070_d

解法

xからKを経由してyに行く最短経路は、xからKへの最短経路とKからyへの最短経路の和になるので、 事前にKから各点への最短経路を求めておけば良い。

最短経路は、対象が閉路の無い木構造なので、幅優先探索や深さ優先探索で求めれば良い。

なお、頂点の数Nが10⁵でありかなり大きいため、隣接情報は隣接行列ではなく、隣接リストで行う必要がある。そうしないと、メモリ制限オーバーとなる。

実装

n = gets.chomp.to_i
edge = Array.new(n+1).map{Array.new}
(n-1).times do
    a, b, c  = gets.chomp.split(" ").map(&:to_i)
    edge[a] << [b, c]
    edge[b] << [a, c]
end

q, k = gets.chomp.split.map(&:to_i)

cost = Array.new(n+1,0)
done = Array.new(n+1,false)

queue = [k]
done[k] = true
cost[k] = 0
while(queue.size > 0)
    v = queue.pop
    
    edge[v].each do |i, j|
        if done[i] == false
            done[i] = true
            cost[i] = cost[v] + j
            queue.push(i)
        else
        end
    end
end

q.times do
    x, y = gets.chomp.split.map(&:to_i)
    puts cost[x] + cost[y]
end