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 Dictionary<long, List<long>> mToNodeListDict = new Dictionary<long, List<long>>();

    class NodeInfo
    {
        internal long MinDFSNo;
        internal long MaxDFSNo;
        internal long EdgeCost;
        internal long EdgeCostSum;
    }

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

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

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

        foreach (string EachStr in InputList.Skip(1).Take((int)N - 1)) {
            SplitAct(EachStr);
            long S = wkArr[0];
            long T = wkArr[1];

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

        long[] CArr = InputList.Last().Split(' ').Select(pX => long.Parse(pX)).ToArray();

        // ノード番号の昇順にソート
        foreach (var EachPair in mToNodeListDict) {
            EachPair.Value.Sort();
        }

        // 状態[ノード]なDP表
        var TreeDPDict = new Dictionary<long, NodeInfo>();

        // 根から葉に向かってDFSしてノードのListを返す
        List<JyoutaiDef1> DFSResult = ExecDFS1(1);
        DFSResult.Reverse(); // 葉から根への順番にする

        // オイラーツアーを行い、DFSでの訪問順序[ノード]なDictを作成
        ExecEulerTour();
        foreach (JyoutaiDef2 EachJyoutai2 in mEulerTourJyoutaiList) {
            if (mDFSNoDict.ContainsKey(EachJyoutai2.Node) == false) {
                mDFSNoDict[EachJyoutai2.Node] = mDFSNoDict.Count + 1;
            }
        }

        foreach (JyoutaiDef1 EachJyoutai in DFSResult) {
            // 親ノードに配る木DP
            long Node = EachJyoutai.CurrNode;
            if (TreeDPDict.ContainsKey(Node) == false) {
                var InitNode = new NodeInfo();
                InitNode.MinDFSNo = mDFSNoDict[Node];
                InitNode.MaxDFSNo = mDFSNoDict[Node];
                InitNode.EdgeCost = CArr[Node - 1];
                InitNode.EdgeCostSum = 0;
                TreeDPDict[Node] = InitNode;
            }

            long ParentNode = EachJyoutai.ParentNode;

            // 根ノードの場合
            if (ParentNode == -1) continue;

            if (TreeDPDict.ContainsKey(ParentNode) == false) {
                NodeInfo InitNode = new NodeInfo();
                InitNode.MinDFSNo = mDFSNoDict[ParentNode];
                InitNode.MaxDFSNo = mDFSNoDict[ParentNode];
                InitNode.MinDFSNo = Math.Min(InitNode.MinDFSNo, mDFSNoDict[Node]);
                InitNode.MaxDFSNo = Math.Max(InitNode.MaxDFSNo, mDFSNoDict[Node]);
                InitNode.EdgeCost = CArr[ParentNode - 1];
                InitNode.EdgeCostSum = 0;
                TreeDPDict[ParentNode] = InitNode;
            }
            TreeDPDict[ParentNode].MinDFSNo =
                Math.Min(TreeDPDict[ParentNode].MinDFSNo, TreeDPDict[Node].MinDFSNo);
            TreeDPDict[ParentNode].MaxDFSNo =
                Math.Max(TreeDPDict[ParentNode].MaxDFSNo, TreeDPDict[Node].MaxDFSNo);
            TreeDPDict[ParentNode].EdgeCost += TreeDPDict[Node].EdgeCost;
            TreeDPDict[ParentNode].EdgeCostSum += TreeDPDict[Node].EdgeCost + TreeDPDict[Node].EdgeCostSum;
        }

        // 木DPの結果
        //foreach (var EachPair in TreeDPDict.OrderBy(pX => pX.Key)) {
        //    Console.WriteLine("TreeDPDict[{0}]=(MinDFSNo={1},MaxDFSNo={2},EdgeCost={3},EdgeCostSum={4})",
        //        EachPair.Key,
        //        EachPair.Value.MinDFSNo,
        //        EachPair.Value.MaxDFSNo,
        //        EachPair.Value.EdgeCost,
        //        EachPair.Value.EdgeCostSum);
        //}

        // DFSNoでのフェニック木
        var Ins_Fenwick_Tree = new Fenwick_Tree(N);
        long Fenwick_Tree_Sum = CArr.Sum();
        for (long I = 0; I <= CArr.GetUpperBound(0); I++) {
            long CurrDFSNo = mDFSNoDict[I + 1];
            long EdgeCost = CArr[I];
            Ins_Fenwick_Tree.Add(CurrDFSNo, EdgeCost);
        }

        var Stk = new Stack<JyoutaiDef3>();
        JyoutaiDef3 WillPush;
        WillPush.CurrNode = 1;
        WillPush.CurrEdgeCostSum = TreeDPDict[1].EdgeCostSum;
        Stk.Push(WillPush);

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

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

            if (mToNodeListDict.ContainsKey(Popped.CurrNode) == false) {
                continue;
            }
            foreach (long EachToNode in mToNodeListDict[Popped.CurrNode]) {
                if (VisitedSet.Add(EachToNode)) {
                    WillPush.CurrNode = EachToNode;
                    WillPush.CurrEdgeCostSum = Popped.CurrEdgeCostSum;
                    long RangeSta = TreeDPDict[EachToNode].MinDFSNo;
                    long RangeEnd = TreeDPDict[EachToNode].MaxDFSNo;
                    long MinusCost = Ins_Fenwick_Tree.GetSum(RangeSta, RangeEnd);
                    long AddCost = Fenwick_Tree_Sum - MinusCost;
                    WillPush.CurrEdgeCostSum += AddCost;
                    WillPush.CurrEdgeCostSum -= MinusCost;
                    Stk.Push(WillPush);
                }
            }
        }
        Console.WriteLine(AnswerList.Min());
    }

    // DFSでの訪問順序[ノード]なDict
    static Dictionary<long, long> mDFSNoDict = new Dictionary<long, long>();

    // 根から葉に向かってDFSしてノードのListを返す
    struct JyoutaiDef1
    {
        internal long CurrNode;
        internal long ParentNode;
    }
    static List<JyoutaiDef1> ExecDFS1(long pRootNode)
    {
        var WillReturn = new List<JyoutaiDef1>();

        var Stk = new Stack<JyoutaiDef1>();
        JyoutaiDef1 WillPush;
        WillPush.CurrNode = pRootNode;
        WillPush.ParentNode = -1;
        Stk.Push(WillPush);

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

        while (Stk.Count > 0) {
            JyoutaiDef1 Popped = Stk.Pop();
            WillReturn.Add(Popped);

            if (mToNodeListDict.ContainsKey(Popped.CurrNode) == false) {
                continue;
            }
            foreach (long EachToNode in mToNodeListDict[Popped.CurrNode]) {
                if (VisitedSet.Add(EachToNode)) {
                    WillPush.CurrNode = EachToNode;
                    WillPush.ParentNode = Popped.CurrNode;
                    Stk.Push(WillPush);
                }
            }
        }
        return WillReturn;
    }

    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()
    {
        var Stk = new Stack<JyoutaiDef2>();
        JyoutaiDef2 WillPush;
        WillPush.Node = 1;
        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(1);

        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);
                    }
                }
            }
        }
    }

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

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

    // コンストラクタ
    internal Fenwick_Tree(long pExternalArrUB)
    {
        mExternalArrUB = pExternalArrUB;

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

    // [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

解説

最初に木DPで、根をノード1とした時のコストを求めます。

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

子ノード達のDFSでの訪問順での最小番号と最大番号は、
最初の木DPで求めてます。