第6回PAST O 最短距離クエリ

問題へのリンク

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("7 6");
            WillReturn.Add("1 2");
            WillReturn.Add("2 3");
            WillReturn.Add("3 4");
            WillReturn.Add("4 5");
            WillReturn.Add("5 6");
            WillReturn.Add("6 7");
            WillReturn.Add("3");
            WillReturn.Add("4 6");
            WillReturn.Add("1 5");
            WillReturn.Add("1 2");
            //2
            //4
            //1
        }
        else if (InputPattern == "Input2") {
            WillReturn.Add("8 8");
            WillReturn.Add("1 6");
            WillReturn.Add("6 7");
            WillReturn.Add("2 5");
            WillReturn.Add("2 3");
            WillReturn.Add("1 8");
            WillReturn.Add("1 5");
            WillReturn.Add("5 6");
            WillReturn.Add("4 8");
            WillReturn.Add("8");
            WillReturn.Add("4 6");
            WillReturn.Add("1 3");
            WillReturn.Add("1 4");
            WillReturn.Add("4 7");
            WillReturn.Add("5 6");
            WillReturn.Add("6 7");
            WillReturn.Add("1 8");
            WillReturn.Add("3 7");
            //3
            //3
            //2
            //4
            //1
            //1
            //1
            //4
        }
        else if (InputPattern == "Input3") {
            WillReturn.Add("11 12");
            WillReturn.Add("6 10");
            WillReturn.Add("9 10");
            WillReturn.Add("8 11");
            WillReturn.Add("3 5");
            WillReturn.Add("1 8");
            WillReturn.Add("7 9");
            WillReturn.Add("2 6");
            WillReturn.Add("3 6");
            WillReturn.Add("4 7");
            WillReturn.Add("3 10");
            WillReturn.Add("1 7");
            WillReturn.Add("1 2");
            WillReturn.Add("6");
            WillReturn.Add("7 10");
            WillReturn.Add("3 8");
            WillReturn.Add("4 7");
            WillReturn.Add("5 9");
            WillReturn.Add("4 5");
            WillReturn.Add("5 10");
            //2
            //4
            //1
            //3
            //5
            //2
        }
        else {
            string wkStr;
            while ((wkStr = Console.ReadLine()) != null) WillReturn.Add(wkStr);
        }
        return WillReturn;
    }

    // 隣接リスト(木の辺)
    static Dictionary<long, List<long>> mToNodeListDict1 = new Dictionary<long, List<long>>();

    // 隣接リスト(木でない辺)
    static Dictionary<long, List<long>> mToNodeListDict2 = new Dictionary<long, List<long>>();

    // 木にあるノードのset
    static HashSet<long> mTreeNodeSet = new HashSet<long>();

    // 最初に登場するInd[ノード]なDict
    static Dictionary<long, long> mFirstApperIndDict = new Dictionary<long, long>();

    // Level[ノード]なDict
    static Dictionary<long, long> mLevelDict = new Dictionary<long, long>();

    // スパーステーブルのコンストラクタに使うLevelのList
    static List<long> mLevelList = new List<long>();

    // 木上の移動候補
    static List<long> mTreeMoveKouhoList = new List<long>();

    static SparseTable<long> mInsSparseTable;

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

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

        SplitAct(InputList[0]);
        long N = wkArr[0];
        long M = wkArr[1];

        var InsUnionFind = new UnionFind();
        for (long I = 1; I <= N; I++) {
            InsUnionFind.MakeSet(I);
        }

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

            long Root1 = InsUnionFind.FindSet(FromNode);
            long Root2 = InsUnionFind.FindSet(ToNode);

            Dictionary<long, List<long>> TargetDict;
            if (Root1 != Root2) {
                InsUnionFind.Unite(Root1, Root2);
                mTreeNodeSet.Add(FromNode);
                mTreeNodeSet.Add(ToNode);
                TargetDict = mToNodeListDict1;
            }
            else {
                TargetDict = mToNodeListDict2;
                mTreeMoveKouhoList.Add(FromNode);
                mTreeMoveKouhoList.Add(ToNode);
            }

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

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

        // 第1引数は、根ノード
        // 第2引数は、親ノード (根ノードの場合は -1
        // 第3引数は、レベル   (根ノードの場合は 1)
        ExecEulerTour(1, -1, 1);

        for (int I = 0; I <= mEulerTourJyoutaiList.Count - 1; I++) {
            int CurrNode = mEulerTourJyoutaiList[I].CurrNode;
            if (mFirstApperIndDict.ContainsKey(CurrNode) == false) {
                mFirstApperIndDict[CurrNode] = I;
            }
            mLevelDict[CurrNode] = mEulerTourJyoutaiList[I].Level;
            mLevelList.Add(mEulerTourJyoutaiList[I].Level);
        }

        mInsSparseTable = new SparseTable<long>(mLevelList, true);

        foreach (string EachStr in InputList.Skip(1 + (int)M + 1)) {
            SplitAct(EachStr);
            long StaNode = wkArr[0];
            long GoalNode = wkArr[1];
            long Result = Dijkstra(StaNode, GoalNode);
            Console.WriteLine(Result);
        }
    }

    // 木上の2点を引数として、距離を返す
    static long DeriveTreeDistance(long pNode1, long pNode2)
    {
        long Ind1 = mFirstApperIndDict[pNode1];
        long Ind2 = mFirstApperIndDict[pNode2];

        long LeftInd = Math.Min(Ind1, Ind2);
        long RightInd = Math.Max(Ind1, Ind2);

        long LCALevel = mInsSparseTable.Query(LeftInd, RightInd);

        long Level1 = mLevelDict[pNode1];
        long Level2 = mLevelDict[pNode2];

        long Answer = 0;
        Answer += Math.Abs(LCALevel - Level1);
        Answer += Math.Abs(LCALevel - Level2);
        return Answer;
    }

    // StaノードとEndノードを引数として、ダイクストラ法を行う
    static long Dijkstra(long pStaNode, long pGoalNode)
    {
        var InsPQueue = new PQueue_Arr();

        // 距離合計[確定ノード]なDict
        var KakuteiNodeDict = new Dictionary<long, long>();
        KakuteiNodeDict.Add(pStaNode, 0);

        // Enqueue処理
        Action<long> EnqueueAct = pCurrNode =>
        {
            // 現在ノードとゴールノードが木にある場合
            if (mTreeNodeSet.Contains(pCurrNode) && mTreeNodeSet.Contains(pGoalNode)) {
                long wkCost = KakuteiNodeDict[pCurrNode] + DeriveTreeDistance(pCurrNode, pGoalNode);
                PQueue_Arr.PQueueJyoutaiDef WillEnqueue;
                WillEnqueue.Node = pGoalNode;
                WillEnqueue.SumCost = wkCost;
                InsPQueue.Enqueue(WillEnqueue);
            }

            // 木上でない辺に隣接したノードに移動
            foreach (long EachToNode in mTreeMoveKouhoList) {
                if (mTreeNodeSet.Contains(pCurrNode) && mTreeNodeSet.Contains(EachToNode)) {
                    long wkCost = KakuteiNodeDict[pCurrNode] + DeriveTreeDistance(pCurrNode, EachToNode);
                    PQueue_Arr.PQueueJyoutaiDef WillEnqueue;
                    WillEnqueue.Node = EachToNode;
                    WillEnqueue.SumCost = wkCost;
                    InsPQueue.Enqueue(WillEnqueue);
                }
            }

            // 木上でない辺を移動
            if (mToNodeListDict2.ContainsKey(pCurrNode)) {
                foreach (long EachToNode in mToNodeListDict2[pCurrNode]) {
                    // 確定ノードならContinue
                    if (KakuteiNodeDict.ContainsKey(EachToNode)) continue;

                    PQueue_Arr.PQueueJyoutaiDef WillEnqueue;
                    WillEnqueue.Node = EachToNode;
                    WillEnqueue.SumCost = KakuteiNodeDict[pCurrNode] + 1;
                    InsPQueue.Enqueue(WillEnqueue);
                }
            }
        };
        EnqueueAct(pStaNode);

        while (InsPQueue.IsEmpty() == false) {
            PQueue_Arr.PQueueJyoutaiDef Dequeued = InsPQueue.Dequeue();

            // 確定ノードならcontinue
            if (KakuteiNodeDict.ContainsKey(Dequeued.Node)) continue;

            // 枝切り
            if (KakuteiNodeDict.ContainsKey(pGoalNode)) break;

            KakuteiNodeDict.Add(Dequeued.Node, Dequeued.SumCost);
            EnqueueAct(Dequeued.Node);
        }

        return KakuteiNodeDict[pGoalNode];
    }

    struct JyoutaiDef_EulerTour
    {
        internal int CurrNode;
        internal int Level;
    }

    // DFSを行い、下記の2通りのタイミングでノードをAddする
    // 1 最初の訪問時
    // 2 子の部分木を訪問時
    static List<JyoutaiDef_EulerTour> mEulerTourJyoutaiList = new List<JyoutaiDef_EulerTour>();
    static void ExecEulerTour(int pCurrNode, int pParentNode, int pLevel)
    {
        JyoutaiDef_EulerTour WillAdd;
        WillAdd.CurrNode = pCurrNode;
        WillAdd.Level = pLevel;
        mEulerTourJyoutaiList.Add(WillAdd);
        if (mToNodeListDict1.ContainsKey(pCurrNode)) {
            foreach (int EachToNode in mToNodeListDict1[pCurrNode]) {
                if (EachToNode == pParentNode) continue;
                ExecEulerTour(EachToNode, pCurrNode, pLevel + 1);
                mEulerTourJyoutaiList.Add(WillAdd);
            }
        }
    }
}

#region PQueue_Arr
// 内部で配列使用の優先度付きキュー
internal class PQueue_Arr
{
    internal struct PQueueJyoutaiDef
    {
        internal long Node;
        internal long SumCost;
    }

    private PQueueJyoutaiDef[] mHeapArr;
    private long mHeapArrCnt = 0;

    // コンストラクタ
    internal PQueue_Arr()
    {
        mHeapArr = new PQueueJyoutaiDef[65535];
    }
    internal bool IsEmpty()
    {
        return mHeapArrCnt == 0;
    }

    // エンキュー処理
    internal void Enqueue(PQueueJyoutaiDef pAddJyoutai)
    {
        long CurrNode = 1 + mHeapArrCnt;
        if (mHeapArr.GetUpperBound(0) < CurrNode) {
            ExtendArr();
        }
        mHeapArr[CurrNode] = pAddJyoutai;
        mHeapArrCnt++;

        while (1 < CurrNode && mHeapArr[CurrNode / 2].SumCost > mHeapArr[CurrNode].SumCost) {
            PQueueJyoutaiDef Swap = mHeapArr[CurrNode];
            mHeapArr[CurrNode] = mHeapArr[CurrNode / 2];
            mHeapArr[CurrNode / 2] = Swap;

            CurrNode /= 2;
        }
    }

    // 配列のExtend
    private void ExtendArr()
    {
        PQueueJyoutaiDef[] NewHeapArr = new PQueueJyoutaiDef[mHeapArrCnt * 2];
        mHeapArr.CopyTo(NewHeapArr, 0);
        mHeapArr = NewHeapArr;
    }

    // デキュー処理
    internal PQueueJyoutaiDef Dequeue()
    {
        PQueueJyoutaiDef TopNode = mHeapArr[1];
        long LastNode = mHeapArrCnt;
        mHeapArr[1] = mHeapArr[LastNode];
        mHeapArrCnt--;

        MinHeapify(1);
        return TopNode;
    }

    // 根ノードを指定し、根から葉へヒープ構築
    private void MinHeapify(long pRootNode)
    {
        if (mHeapArrCnt <= 1) {
            return;
        }

        long Left = pRootNode * 2;
        long Right = pRootNode * 2 + 1;

        // 左の子、自分、右の子で値が最小のノードを選ぶ
        long Smallest = mHeapArr[pRootNode].SumCost;
        long SmallestNode = pRootNode;

        if (Left <= mHeapArrCnt && mHeapArr[Left].SumCost < Smallest) {
            Smallest = mHeapArr[Left].SumCost;
            SmallestNode = Left;
        }
        if (Right <= mHeapArrCnt && mHeapArr[Right].SumCost < Smallest) {
            Smallest = mHeapArr[Right].SumCost;
            SmallestNode = Right;
        }

        // 子ノードのほうが大きい場合
        if (SmallestNode != pRootNode) {
            PQueueJyoutaiDef Swap = mHeapArr[SmallestNode];
            mHeapArr[SmallestNode] = mHeapArr[pRootNode];
            mHeapArr[pRootNode] = Swap;

            // 再帰的に呼び出し
            MinHeapify(SmallestNode);
        }
    }
}
#endregion

#region UnionFind
// UnionFindクラス
internal class UnionFind
{
    private class NodeInfoDef
    {
        internal long ParentNode;
        internal long Rank;
    }
    private Dictionary<long, NodeInfoDef> mNodeInfoDict =
        new Dictionary<long, NodeInfoDef>();

    // 要素が1つである木を森に追加
    internal void MakeSet(long pNode)
    {
        NodeInfoDef WillAdd = new NodeInfoDef();
        WillAdd.ParentNode = pNode;
        WillAdd.Rank = 0;
        mNodeInfoDict[pNode] = WillAdd;
    }

    // 合併処理
    internal void Unite(long pX, long pY)
    {
        long XNode = FindSet(pX);
        long YNode = FindSet(pY);
        long XRank = mNodeInfoDict[XNode].Rank;
        long YRank = mNodeInfoDict[YNode].Rank;

        if (XRank > YRank) {
            mNodeInfoDict[YNode].ParentNode = XNode;
        }
        else {
            mNodeInfoDict[XNode].ParentNode = YNode;
            if (XRank == YRank) {
                mNodeInfoDict[YNode].Rank++;
            }
        }
    }

    // ノードを引数として、木の根を取得
    internal long FindSet(long pTargetNode)
    {
        // 根までの経路上のノードのList
        var PathNodeList = new List<long>();

        long CurrNode = pTargetNode;
        while (CurrNode != mNodeInfoDict[CurrNode].ParentNode) {
            PathNodeList.Add(CurrNode);
            CurrNode = mNodeInfoDict[CurrNode].ParentNode;
        }

        // 経路圧縮 (親ポインタの付け替え)
        foreach (long EachPathNode in PathNodeList) {
            mNodeInfoDict[EachPathNode].ParentNode = CurrNode;
        }
        return CurrNode;
    }

    internal void DebugPrint()
    {
        foreach (var EachPair in mNodeInfoDict.OrderBy(pX => pX.Key)) {
            Console.WriteLine("mNodeInfoDict[{0}].ParentNode={1}",
                EachPair.Key, EachPair.Value.ParentNode);
        }
    }
}
#endregion

#region SparseTable
// スパーステーブル
internal class SparseTable<Type>
{
    private Type[] mInitArr;
    private long UB_0;
    private long UB_1;

    // 最小値か最大値[開始Ind,2の指数]なArr
    private Type[,] mMinOrMaxArr;

    // Log2の値[2べきな値] なDict
    private Dictionary<long, long> mLogDict = new Dictionary<long, long>();

    // 最小値を取得するか？
    private bool mIsMin = false;

    // コンストラクタ
    internal SparseTable(IEnumerable<Type> pInitEnum, bool pIsMin)
    {
        mIsMin = pIsMin;

        mInitArr = pInitEnum.ToArray();
        UB_0 = mInitArr.GetUpperBound(0);

        long Sisuu = 0;
        long Beki2 = 1;
        while (true) {
            mLogDict[Beki2] = Sisuu;
            if (Beki2 > mInitArr.Length) {
                break;
            }
            Beki2 *= 2;
            Sisuu++;
        }
        UB_1 = Sisuu;

        mMinOrMaxArr = new Type[UB_0 + 1, UB_1 + 1];

        // 長さ1の分を初期化
        for (long I = 0; I <= UB_0; I++) {
            mMinOrMaxArr[I, 0] = mInitArr[I];
        }

        for (long LoopLength = 2; LoopLength < long.MaxValue; LoopLength *= 2) {
            bool WillBreak = true;
            long HalfRange = LoopLength / 2;
            for (long I = 0; I <= UB_0; I++) {
                long StaInd1 = I;
                long EndInd1 = I + HalfRange - 1;
                long StaInd2 = EndInd1 + 1;
                long EndInd2 = StaInd2 + HalfRange - 1;

                if (EndInd2 > UB_0) break;

                var KouhoList = new List<Type>();
                KouhoList.Add(mMinOrMaxArr[I, mLogDict[HalfRange]]);
                KouhoList.Add(mMinOrMaxArr[StaInd2, mLogDict[HalfRange]]);

                if (mIsMin) {
                    mMinOrMaxArr[I, mLogDict[LoopLength]] = KouhoList.Min();
                }
                else {
                    mMinOrMaxArr[I, mLogDict[LoopLength]] = KouhoList.Max();
                }

                WillBreak = false;
            }
            if (WillBreak) break;
        }
    }

    // 閉区間 [L,R]の最小値または最大値を求める
    internal Type Query(long pL, long pR)
    {
        // 区間を内包する2べきを返す
        long Beki2 = 1;
        long Sisuu = 0;
        while (true) {
            long LeftRangeMax = pL + Beki2 - 1;
            long RightRangeMin = pR - Beki2 + 1;

            if (LeftRangeMax + 1 >= RightRangeMin) {
                break;
            }
            Beki2 *= 2;
            Sisuu++;
        }

        var KouhoList = new List<Type>();
        KouhoList.Add(mMinOrMaxArr[pL, Sisuu]);
        KouhoList.Add(mMinOrMaxArr[pR - Beki2 + 1, Sisuu]);

        if (mIsMin) {
            return KouhoList.Min();
        }
        return KouhoList.Max();
    }
}
#endregion

解説

木上での移動は、LCA経由で距離を求める。
木上でない移動は、ナイーブにコスト計算とし、
ダイクストラ法を使ってます。