ABC-070-D Transit Tree Path

■■■問題■■■

N頂点の木が与えられます。

木とはグラフの一種であり、頂点の数をNとすると、辺の数が N-1 本である閉路のない連結グラフです。
i(1 <= i <= N-1) 番目の辺は、頂点aiと頂点biを距離ciで結びます。

また、Q個の質問クエリと整数Kが与えられます。
●j(1 <= j <= Q) 番目の質問クエリでは、頂点xjから頂点Kを経由しつつ、
  頂点yjまで移動する場合の最短経路の距離を求めてください。

■■■入力■■■

N
a1 b1 c1
・
・
・
aN-1 bN-1 cN-1
Q K
x1 y1
・
・
・
xQ yQ

●3 <= N <= 10万
●1 <= ai,bi <= N(1 <= i <= N-1)
●1 <= ci <= 10億(1 <= i <= N-1)
●与えられるグラフは木である
●1 <= Q <= 10万
●1 <= K <= N
●1 <= xj,yj <= N(1 <= j <= Q)
●xj != yj(1 <= j <= Q)
●xj != K,yj != K(1 <= j <= Q)

■■■出力■■■

質問クエリの解答をQ行出力せよ。
j(1 <= j <= Q) 行目には、j番目のクエリの答えを出力せよ。

C#のソース

using System;
using System.Collections.Generic;
using System.Linq;

class Program
{
    static string InputPattern = "Input1";

    static List<string> GetInputList()
    {
        var WillReturn = new List<string>();

        if (InputPattern == "Input1") {
            WillReturn.Add("5");
            WillReturn.Add("1 2 1");
            WillReturn.Add("1 3 1");
            WillReturn.Add("2 4 1");
            WillReturn.Add("3 5 1");
            WillReturn.Add("3 1");
            WillReturn.Add("2 4");
            WillReturn.Add("2 3");
            WillReturn.Add("4 5");
            //3
            //2
            //4
            //与えられた3つの質問クエリに対する最短経路は以下の通りです。
            //●1つ目の質問クエリ: 頂点2 → 頂点1 → 頂点2 → 頂点4 : 距離 1+1+1=3
            //●2つ目の質問クエリ: 頂点2 → 頂点1 → 頂点3 : 距離 1+1=2
            //●3つ目の質問クエリ: 頂点4 → 頂点2 → 頂点1 → 頂点3 → 頂点5 : 距離 1+1+1+1=4
        }
        else if (InputPattern == "Input2") {
            WillReturn.Add("7");
            WillReturn.Add("1 2 1");
            WillReturn.Add("1 3 3");
            WillReturn.Add("1 4 5");
            WillReturn.Add("1 5 7");
            WillReturn.Add("1 6 9");
            WillReturn.Add("1 7 11");
            WillReturn.Add("3 2");
            WillReturn.Add("1 3");
            WillReturn.Add("4 5");
            WillReturn.Add("6 7");
            //5
            //14
            //22
            //質問クエリに対する最短経路は、必ず頂点 K=2 を通過する必要があります。
        }
        else if (InputPattern == "Input3") {
            WillReturn.Add("10");
            WillReturn.Add("1 2 1000000000");
            WillReturn.Add("2 3 1000000000");
            WillReturn.Add("3 4 1000000000");
            WillReturn.Add("4 5 1000000000");
            WillReturn.Add("5 6 1000000000");
            WillReturn.Add("6 7 1000000000");
            WillReturn.Add("7 8 1000000000");
            WillReturn.Add("8 9 1000000000");
            WillReturn.Add("9 10 1000000000");
            WillReturn.Add("1 1");
            WillReturn.Add("9 10");
            //17000000000
        }
        else {
            string wkStr;
            while ((wkStr = Console.ReadLine()) != null) WillReturn.Add(wkStr);
        }
        return WillReturn;
    }

    static int N;

    //辺の情報
    struct EdgeInfoDef
    {
        internal int ToPos;
        internal long Cost;
    }
    //隣接リストで、辺の情報のList[頂点番号]を管理
    static Dictionary<int, List<EdgeInfoDef>> mEdgeInfoDict =
        new Dictionary<int, List<EdgeInfoDef>>();

    static int K;

    //Kからの距離[頂点番号]な配列
    static long[] mKyoriArr;

    static void Main()
    {
        List<string> InputList = GetInputList();

        int[] wkArr = { };
        Action<string> SplitAct = pStr =>
            wkArr = pStr.Split(' ').Select(X => int.Parse(X)).ToArray();

        N = int.Parse(InputList[0]);

        for (int I = 1; I <= N - 1; I++) {
            SplitAct(InputList[I]);
            int A = wkArr[0];
            int B = wkArr[1];
            int C = wkArr[2];

            if (mEdgeInfoDict.ContainsKey(A) == false)
                mEdgeInfoDict.Add(A, new List<EdgeInfoDef>());
            if (mEdgeInfoDict.ContainsKey(B) == false)
                mEdgeInfoDict.Add(B, new List<EdgeInfoDef>());

            mEdgeInfoDict[A].Add(new EdgeInfoDef() { ToPos = B, Cost = C });
            mEdgeInfoDict[B].Add(new EdgeInfoDef() { ToPos = A, Cost = C });
        }

        SplitAct(InputList[N]);
        K = wkArr[1];

        ExecDFS();

        Console.WriteLine("K({0})から各頂点までの距離", K);
        for (int I = 1; I <= N; I++) {
            Console.WriteLine("mKyoriArr[{0}]={1}", I, mKyoriArr[I]);
        }

        //問い合わせの処理
        for (int I = N + 1; I <= InputList.Count - 1; I++) {
            SplitAct(InputList[I]);
            int X = wkArr[0];
            int Y = wkArr[1];

            long Answer = 0;
            Answer += mKyoriArr[X];
            Answer += mKyoriArr[Y];
            Console.WriteLine(Answer);
        }
    }

    struct JyoutaiDef
    {
        internal int CurrPos;
        internal long CostSum;
    }

    //深さ優先探索でKから各頂点までの距離を求める
    //距離[頂点番号]な配列を求める
    static void ExecDFS()
    {
        mKyoriArr = new long[N + 1];

        var Stk = new Stack<JyoutaiDef>();
        JyoutaiDef WillPush;
        WillPush.CurrPos = K;
        WillPush.CostSum = 0;
        mKyoriArr[K] = 0;
        Stk.Push(WillPush);

        var VisitedSet = new HashSet<int>();
        VisitedSet.Add(K);

        while (Stk.Count > 0) {
            JyoutaiDef Popped = Stk.Pop();

            mKyoriArr[Popped.CurrPos] = Popped.CostSum;

            foreach (EdgeInfoDef EachEdgeInfo in mEdgeInfoDict[Popped.CurrPos]) {
                WillPush.CurrPos = EachEdgeInfo.ToPos;
                WillPush.CostSum = Popped.CostSum + EachEdgeInfo.Cost;

                //再訪禁止
                if (VisitedSet.Add(WillPush.CurrPos) == false)
                    continue;

                Stk.Push(WillPush);
            }
        }
    }
}

デバッグ情報付の実行結果

K(1)から各頂点までの距離
mKyoriArr[1]=0
mKyoriArr[2]=1
mKyoriArr[3]=1
mKyoriArr[4]=2
mKyoriArr[5]=2
3
2
4

解説

隣接行列ではなく、隣接リストで辺の情報を管理して、
頂点Kから各頂点までの最短距離を深さ優先探索で求めます。