ABC267-F Exactly K Steps

問題へのリンク

C#のソース

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

class Program
{
    static string InputPattern = "InputX";

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

        if (InputPattern == "Input1") {
            WillReturn.Add("5");
            WillReturn.Add("1 2");
            WillReturn.Add("2 3");
            WillReturn.Add("3 4");
            WillReturn.Add("3 5");
            WillReturn.Add("3");
            WillReturn.Add("2 2");
            WillReturn.Add("5 3");
            WillReturn.Add("3 3");
            //4
            //1
            //-1
        }
        else if (InputPattern == "Input2") {
            WillReturn.Add("10");
            WillReturn.Add("1 2");
            WillReturn.Add("2 3");
            WillReturn.Add("3 5");
            WillReturn.Add("2 8");
            WillReturn.Add("3 4");
            WillReturn.Add("4 6");
            WillReturn.Add("4 9");
            WillReturn.Add("5 7");
            WillReturn.Add("9 10");
            WillReturn.Add("5");
            WillReturn.Add("1 1");
            WillReturn.Add("2 2");
            WillReturn.Add("3 3");
            WillReturn.Add("4 4");
            WillReturn.Add("5 5");
            //2
            //4
            //10
            //-1
            //-1
        }
        else {
            string wkStr;
            while ((wkStr = Console.ReadLine()) != null) WillReturn.Add(wkStr);
        }
        return WillReturn;
    }

    static int mN;

    // 隣接リスト
    static Dictionary<int, List<int>> mToNodeListDict = new Dictionary<int, List<int>>();

    class QueryInfoDef
    {
        internal int Time;
        internal int K;
        internal int Answer;
    }
    static Dictionary<int, List<QueryInfoDef>> mQueryInfoListDict =
        new Dictionary<int, List<QueryInfoDef>>();

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

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

        SplitAct(InputList[0]);
        mN = wkArr[0];

        foreach (string EachStr in InputList.Skip(1).Take(mN - 1)) {
            SplitAct(EachStr);
            int FromNode = wkArr[0];
            int ToNode = wkArr[1];

            if (mToNodeListDict.ContainsKey(FromNode) == false) {
                mToNodeListDict[FromNode] = new List<int>();
            }
            if (mToNodeListDict.ContainsKey(ToNode) == false) {
                mToNodeListDict[ToNode] = new List<int>();
            }
            mToNodeListDict[FromNode].Add(ToNode);
            mToNodeListDict[ToNode].Add(FromNode);
        }

        // クエリを先読みしておく
        int CurrTime = 0;
        foreach (string EachStr in InputList.Skip(1 + mN - 1 + 1)) {
            SplitAct(EachStr);
            int U = wkArr[0];
            int K = wkArr[1];
            if (mQueryInfoListDict.ContainsKey(U) == false) {
                mQueryInfoListDict[U] = new List<QueryInfoDef>();
            }
            var WillAdd = new QueryInfoDef();
            WillAdd.Time = CurrTime++;
            WillAdd.K = K;
            WillAdd.Answer = -1;
            mQueryInfoListDict[U].Add(WillAdd);
        }

        // 直径をなす2つの葉ノードを求める
        int Leaf1 = ExecBFS(1);
        int Leaf2 = ExecBFS(Leaf1);

        var AncestorList = new List<int>();
        Action ForEachAct = () =>
        {
            foreach (JyoutaiDef2 EachJyoutai2 in mEulerTourJyoutaiList) {
                if (EachJyoutai2.IsStart) {
                    AncestorList.Add(EachJyoutai2.Node);

                    if (mQueryInfoListDict.ContainsKey(EachJyoutai2.Node)) {
                        for (int I = 0; I <= mQueryInfoListDict[EachJyoutai2.Node].Count - 1; I++) {
                            int CurrUB = AncestorList.Count - 1;
                            int AnswerInd = CurrUB - mQueryInfoListDict[EachJyoutai2.Node][I].K;
                            if (AnswerInd >= 0) {
                                mQueryInfoListDict[EachJyoutai2.Node][I].Answer = AncestorList[AnswerInd];
                            }
                        }
                    }
                }
                if (EachJyoutai2.IsEnd) {
                    AncestorList.RemoveAt(AncestorList.Count - 1);
                }
            }
        };

        // オイラーツアー1を行う
        mEulerTourJyoutaiList.Clear();
        ExecEulerTour(Leaf1);
        ForEachAct();

        // オイラーツアー2を行う
        mEulerTourJyoutaiList.Clear();
        ExecEulerTour(Leaf2);
        ForEachAct();

        var AllQueryList = new List<QueryInfoDef>();
        foreach (var EachPair in mQueryInfoListDict) {
            AllQueryList.AddRange(EachPair.Value);
        }
        var sb = new System.Text.StringBuilder();
        foreach (QueryInfoDef EachQueryInfo in AllQueryList.OrderBy(pX => pX.Time)) {
            sb.Append(EachQueryInfo.Answer);
            sb.AppendLine();
        }
        Console.Write(sb.ToString());
    }

    struct JyoutaiDef1
    {
        internal int CurrNode;
    }

    // 開始ノードを引数として、一番遠い葉ノードを返す
    static int ExecBFS(int pStaNode)
    {
        var Que = new Queue<JyoutaiDef1>();
        JyoutaiDef1 WillEnqueue;
        WillEnqueue.CurrNode = pStaNode;
        Que.Enqueue(WillEnqueue);

        var VisitedSet = new HashSet<int>();
        VisitedSet.Add(pStaNode);
        int AnswerNode = pStaNode;
        while (Que.Count > 0) {
            JyoutaiDef1 Dequeued = Que.Dequeue();
            AnswerNode = Dequeued.CurrNode;

            if (mToNodeListDict.ContainsKey(Dequeued.CurrNode)) {
                foreach (int EachToNode in mToNodeListDict[Dequeued.CurrNode]) {
                    if (VisitedSet.Add(EachToNode)) {
                        WillEnqueue.CurrNode = EachToNode;
                        Que.Enqueue(WillEnqueue);
                    }
                }
            }
        }
        return AnswerNode;
    }

    struct JyoutaiDef2
    {
        internal int Node;
        internal int Level;
        internal bool IsStart;
        internal bool IsEnd;
    }

    // オイラーツアーを行い、下記の2通りのタイミングでノードをAddする
    // 1 最初の訪問時
    // 2 子の部分木を訪問完了時
    static List<JyoutaiDef2> mEulerTourJyoutaiList = new List<JyoutaiDef2>();
    static void ExecEulerTour(int pStaNode)
    {
        var Stk = new Stack<JyoutaiDef2>();
        JyoutaiDef2 WillPush;
        WillPush.Node = pStaNode;
        WillPush.Level = 0;

        WillPush.IsStart = false; WillPush.IsEnd = true;
        Stk.Push(WillPush);
        WillPush.IsStart = true; WillPush.IsEnd = false;
        Stk.Push(WillPush);

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

        while (Stk.Count > 0) {
            JyoutaiDef2 Popped = Stk.Pop();
            mEulerTourJyoutaiList.Add(Popped);
            if (Popped.IsEnd) {
                continue;
            }
            if (mToNodeListDict.ContainsKey(Popped.Node)) {
                foreach (int EachToNode in mToNodeListDict[Popped.Node]) {
                    if (VisitedSet.Add(EachToNode)) {
                        WillPush.Node = EachToNode;
                        WillPush.Level = Popped.Level + 1;

                        WillPush.IsStart = false; WillPush.IsEnd = true;
                        Stk.Push(WillPush);
                        WillPush.IsStart = true; WillPush.IsEnd = false;
                        Stk.Push(WillPush);
                    }
                }
            }
        }
    }
}

解説

最初に木の直径での2つの葉ノードを求めます。
次に、上記の2つの葉ノードからオイラーツアーを行えば、クエリの解を設定できます。