ABC348-E Minimize Sum of Distances

問題へのリンク

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

    static long[] GetSplitArr(string pStr)
    {
        return (pStr == "" ? new string[0] : pStr.Split(' ')).Select(pX => long.Parse(pX)).ToArray();
    }

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

    // 根付き木の解析クラス
    static List<ClassTreeAnalyze.NodeInfoDef> mNodeInfoList;
    static Dictionary<long, ClassTreeAnalyze.NodeInfoDef> mCurrNodeDict;
    static Dictionary<long, ClassTreeAnalyze.NodeInfoDef> mDFSNoCurrDict;

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

        long[] wkArr = { };
        Action<string> SplitAct = (pStr) => wkArr = GetSplitArr(pStr);

        int NodeCnt = int.Parse(InputList[0]);
        SplitAct(InputList.Last());
        long[] CArr = wkArr;

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

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

        // 根付き木の解析クラス
        ClassTreeAnalyze.GetTreeAnalyzeResult(1, mToNodeListDict,
            out mNodeInfoList, out mCurrNodeDict, out mDFSNoCurrDict);

        // DFSNoでのフェニック木
        var Ins_Fenwick_Tree = new Fenwick_Tree(NodeCnt);
        long Fenwick_Tree_Sum = CArr.Sum();
        for (long I = 0; I <= CArr.GetUpperBound(0); I++) {
            long DFSNoCurr = mCurrNodeDict[I + 1].DFSNoCurr;
            Ins_Fenwick_Tree[DFSNoCurr] = CArr[I];
        }

        // 1を根とした時のコスト合計を求める
        long RootNode1CostSum = 0;
        foreach (ClassTreeAnalyze.NodeInfoDef EachTreeDFSNodeInfo in mNodeInfoList) {
            RootNode1CostSum += CArr[EachTreeDFSNodeInfo.CurrNode - 1] * (EachTreeDFSNodeInfo.Level - 1);
        }

        var Stk = new Stack<JyoutaiDef>();
        JyoutaiDef WillPush;
        WillPush.CurrNode = 1;
        WillPush.CurrEdgeCostSum = RootNode1CostSum;
        Stk.Push(WillPush);

        var VisitedSet = new HashSet<long>();
        VisitedSet.Add(1);

        var AnswerList = new List<long>();
        while (Stk.Count > 0) {
            JyoutaiDef Popped = Stk.Pop();
            AnswerList.Add(Popped.CurrEdgeCostSum);

            foreach (long EachToNode in mToNodeListDict[Popped.CurrNode]) {
                if (VisitedSet.Add(EachToNode)) {
                    WillPush.CurrNode = EachToNode;
                    WillPush.CurrEdgeCostSum = Popped.CurrEdgeCostSum;
                    long RangeSta = mCurrNodeDict[EachToNode].DFSNoSta;
                    long RangeEnd = mCurrNodeDict[EachToNode].DFSNoEnd;
                    long MinusCost = Ins_Fenwick_Tree.GetSum(RangeSta, RangeEnd);
                    long PlusCost = Fenwick_Tree_Sum - MinusCost;
                    WillPush.CurrEdgeCostSum += PlusCost;
                    WillPush.CurrEdgeCostSum -= MinusCost;
                    Stk.Push(WillPush);
                }
            }
        }
        Console.WriteLine(AnswerList.Min());
    }

    // 根から葉に向かってDFS
    struct JyoutaiDef
    {
        internal long CurrNode;
        internal long CurrEdgeCostSum;
    }
}

#region ClassTreeAnalyze
// 根付き木の解析クラス
internal class ClassTreeAnalyze
{
    const long NonExistsParent = -1; // 存在しない親ノード

    // ノード情報
    internal class NodeInfoDef
    {
        internal long CurrNode;   // 現在ノード
        internal long ParentNode; // 親ノード
        internal long Level;      // レベル
        internal long DFSNoCurr;  // DFSでの訪問順
        internal long DFSNoSta;   // DFSでの訪問順での子孫ノードの区間開始
        internal long DFSNoEnd;   // DFSでの訪問順での子孫ノードの区間終了
    }

    // ルートノードと連結リストを引数とし、下記を返す
    // ノード情報のList
    // ノード情報[現在ノード]なDict
    // ノード情報[DFSでの訪問順]
    static internal void GetTreeAnalyzeResult(long pRootNode, Dictionary<long, List<long>> pToNodeListDict,
        out List<ClassTreeAnalyze.NodeInfoDef> pNodeInfoList,
        out Dictionary<long, ClassTreeAnalyze.NodeInfoDef> pCurrNodeDict,
        out Dictionary<long, ClassTreeAnalyze.NodeInfoDef> pDFSNoCurrDict)
    {
        // ノード情報のList
        pNodeInfoList = new List<NodeInfoDef>();

        // ノード情報[現在ノード]なDict
        pCurrNodeDict = new Dictionary<long, NodeInfoDef>();

        // ノード情報[DFSでの訪問順]なDict
        pDFSNoCurrDict = new Dictionary<long, NodeInfoDef>();

        // DFSする
        var Stk = new Stack<JyoutaiDefTreeDFS>();
        JyoutaiDefTreeDFS WillPush;
        WillPush.CurrNode = pRootNode;
        WillPush.ParentNode = NonExistsParent; // 親ノードなし
        WillPush.Level = 1;
        Stk.Push(WillPush);

        var VisitedSet = new HashSet<long>();
        VisitedSet.Add(pRootNode);
        while (Stk.Count > 0) {
            JyoutaiDefTreeDFS Popped = Stk.Pop();

            var WillAdd = new NodeInfoDef();
            WillAdd.CurrNode = Popped.CurrNode;
            WillAdd.ParentNode = Popped.ParentNode;
            WillAdd.Level = Popped.Level;
            WillAdd.DFSNoCurr = pNodeInfoList.Count + 1;
            WillAdd.DFSNoSta = WillAdd.DFSNoCurr;
            WillAdd.DFSNoEnd = WillAdd.DFSNoCurr;
            pNodeInfoList.Add(WillAdd);

            pCurrNodeDict[Popped.CurrNode] = WillAdd;
            pDFSNoCurrDict[WillAdd.DFSNoCurr] = WillAdd;

            foreach (long EachToNode in pToNodeListDict[Popped.CurrNode].OrderByDescending(pX => pX)) {
                if (VisitedSet.Add(EachToNode)) {
                    WillPush.CurrNode = EachToNode;
                    WillPush.ParentNode = Popped.CurrNode;
                    WillPush.Level = Popped.Level + 1;
                    Stk.Push(WillPush);
                }
            }
        }

        // 葉から根に木DP
        for (int I = pNodeInfoList.Count - 1; 0 <= I; I--) {
            long CurrNode = pNodeInfoList[I].CurrNode;
            long ParentNode = pNodeInfoList[I].ParentNode;

            // 親ノードなしの場合
            if (ParentNode == NonExistsParent) continue;

            long ParentDFSNoMin = pCurrNodeDict[ParentNode].DFSNoSta;
            pCurrNodeDict[ParentNode].DFSNoSta = Math.Min(ParentDFSNoMin, pCurrNodeDict[CurrNode].DFSNoSta);

            long ParentDFSNoMax = pCurrNodeDict[ParentNode].DFSNoEnd;
            pCurrNodeDict[ParentNode].DFSNoEnd = Math.Max(ParentDFSNoMax, pCurrNodeDict[CurrNode].DFSNoEnd);
        }

        //DebugPrint(pTreeDFSNodeInfoList);
    }

    private struct JyoutaiDefTreeDFS
    {
        internal long CurrNode;
        internal long ParentNode;
        internal long Level;
    }

    // デバッグ出力
    static void DebugPrint(List<NodeInfoDef> pTreeDFSNodeInfoList)
    {
        foreach (NodeInfoDef EachTreeDFSNodeInfo in pTreeDFSNodeInfoList) {
            Console.WriteLine(
                "CurrNode={0,2},ParentNode={1,2},Level={2,2},DFSNoCurr={3,2},DFSNoSta={4,2},DFSNoEnd={5,2}",
                EachTreeDFSNodeInfo.CurrNode, EachTreeDFSNodeInfo.ParentNode, EachTreeDFSNodeInfo.Level,
                EachTreeDFSNodeInfo.DFSNoCurr, EachTreeDFSNodeInfo.DFSNoSta, EachTreeDFSNodeInfo.DFSNoEnd);
        }
    }
}
#endregion

// フェニック木
#region Fenwick_Tree
internal class Fenwick_Tree
{
    private long[] mBitArr;
    private long mExternalArrUB;

    // ノードのIndexの列挙を返す
    internal IEnumerable<long> GetNodeIndEnum()
    {
        for (long I = 0; I <= mExternalArrUB; I++) {
            yield return I;
        }
    }

    // 木のノードのUBを返す
    internal long GetUB()
    {
        return mExternalArrUB;
    }

    // コンストラクタ(外部配列のUBのみ指定)
    internal Fenwick_Tree(long pExternalArrUB)
    {
        mExternalArrUB = pExternalArrUB;

        // フェニック木の外部配列は0オリジンで、
        // フェニック木の内部配列は1オリジンなため、2を足す
        mBitArr = new long[pExternalArrUB + 2];
    }

    // コンストラクタ(初期化用の配列指定)
    internal Fenwick_Tree(long[] pArr)
        : this(pArr.GetUpperBound(0))
    {
        for (long I = 0; I <= pArr.GetUpperBound(0); I++) {
            this.Add(I, pArr[I]);
        }
    }

    // コンストラクタ(初期化用のList指定)
    internal Fenwick_Tree(List<long> pList)
        : this(pList.Count - 1)
    {
        for (int I = 0; I <= pList.Count - 1; I++) {
            this.Add(I, pList[I]);
        }
    }

    // インデクサ
    internal long this[long pInd]
    {
        get { return GetSum(pInd, pInd); }
        set { Add(pInd, value - GetSum(pInd, pInd)); }
    }

    // [pSta,pEnd] のSumを返す
    internal long GetSum(long pSta, long pEnd)
    {
        return GetSum(pEnd) - GetSum(pSta - 1);
    }

    // [0,pEnd] のSumを返す
    internal long GetSum(long pEnd)
    {
        pEnd++; // 1オリジンに変更

        long Sum = 0;
        while (pEnd >= 1) {
            Sum += mBitArr[pEnd];
            pEnd -= pEnd & -pEnd;
        }
        return Sum;
    }

    // [I] に Xを加算
    internal void Add(long pI, long pX)
    {
        pI++; // 1オリジンに変更

        while (pI <= mBitArr.GetUpperBound(0)) {
            mBitArr[pI] += pX;
            pI += pI & -pI;
        }
    }
}
#endregion

解説

最初に根付き木の解析クラスで、
子ノード達のDFSでの訪問順での
{親ノード、レベル、DFSでの訪問順、
 DFSでの訪問順での子孫ノードの区間開始、
 DFSでの訪問順での子孫ノードの区間終了}
を求めます。

次に、根をノード1とした時のコストを求めます。

それから、DFSでReRootingしますが、
コストの増減は、DFSでの訪問順を考えれば、
フェニック木の区間和から高速に求めることができます。