ABC222-F Expensive Expense

問題へのリンク

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("3");
            WillReturn.Add("1 2 2");
            WillReturn.Add("2 3 3");
            WillReturn.Add("1 2 3");
            //8
            //6
            //6
        }
        else if (InputPattern == "Input2") {
            WillReturn.Add("6");
            WillReturn.Add("1 2 3");
            WillReturn.Add("1 3 1");
            WillReturn.Add("1 4 4");
            WillReturn.Add("1 5 1");
            WillReturn.Add("1 6 5");
            WillReturn.Add("9 2 6 5 3 100");
            //105
            //108
            //106
            //109
            //106
            //14
        }
        else if (InputPattern == "Input3") {
            WillReturn.Add("6");
            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("1 2 3 4 5 6");
            //5000000006
            //4000000006
            //3000000006
            //3000000001
            //4000000001
            //5000000001
        }
        else {
            string wkStr;
            while ((wkStr = Console.ReadLine()) != null) WillReturn.Add(wkStr);
        }
        return WillReturn;
    }

    static List<string> mInputList = new List<string>();

    static void Main()
    {
        mInputList = GetInputList();
        Solve();
    }

    static long mN;

    struct EdgeInfoDef
    {
        internal long ToNode;
        internal long Cost;
    }
    static Dictionary<long, List<EdgeInfoDef>> mEdgeInfoListDict = new Dictionary<long, List<EdgeInfoDef>>();

    static long[] mDArr;

    // 木DPとReRootingの設定01 配られてる値が無い時の値
    static long DeriveDefaultNodeVal()
    {
        return 0;
    }

    // 木DPとReRootingの設定02 ノード値を引数として、配る値を返す
    static long DeriveSendVal(long pNodeVal, long pFromNode, long pToNode)
    {
        string EdgeHash = DeriveEdgeHash(pFromNode, pToNode);
        long SendVal = pNodeVal + mEdgeCostDict[EdgeHash];
        return SendVal;
    }

    // 木DPとReRootingの設定03 ノード値と、配られた値を、引数とし、新しいノード値を返す
    static long DeriveNewNodeVal(long pNodeVal, long pSendVal)
    {
        return Math.Max(pNodeVal, pSendVal);
    }

    // 木DPとReRootingの設定04 配る値のListを引数として、最終的な配る値を返す
    static long DeriveMergedSendVal(List<long> pSendValList)
    {
        long SendVal = DeriveDefaultNodeVal();
        foreach (long EachSendVal in pSendValList) {
            SendVal = DeriveNewNodeVal(SendVal, EachSendVal);
        }
        return SendVal;
    }

    static void Solve()
    {
        mN = long.Parse(mInputList[0]);

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

        Action<long, long, long> AddEdge = (pFrom, pTo, pCost) =>
        {
            if (mEdgeInfoListDict.ContainsKey(pFrom) == false) {
                mEdgeInfoListDict[pFrom] = new List<EdgeInfoDef>();
            }
            EdgeInfoDef WillAdd;
            WillAdd.ToNode = pTo;
            WillAdd.Cost = pCost;
            mEdgeInfoListDict[pFrom].Add(WillAdd);
        };

        foreach (string EachStr in mInputList.Skip(1).Take((int)mN - 1)) {
            SplitAct(EachStr);
            long A = wkArr[0];
            long B = wkArr[1];
            long C = wkArr[2];

            AddEdge(A, B, C);
            AddEdge(B, A, C);
        }

        mDArr = mInputList.Last().Split(' ').Select(pX => long.Parse(pX)).ToArray();

        // 滞在コストは、葉を作り、ノードでなく枝がコストを持つようにする
        for (long I = 0; I <= mDArr.GetUpperBound(0); I++) {
            AddEdge(I + 1, I + 1000000, mDArr[I]);
            AddEdge(I + 1000000, I + 1, mDArr[I]);
        }

        // 前後に累積和を使うので、ノード番号の昇順にソート
        foreach (var EachPair in mEdgeInfoListDict) {
            EachPair.Value.Sort((A, B) => A.ToNode.CompareTo(B.ToNode));
        }

        // 枝のコストのDictを作っておく
        foreach (var EachPair in mEdgeInfoListDict) {
            foreach (EdgeInfoDef EachEdgeInfo in EachPair.Value) {
                long FromNode = EachPair.Key;
                long ToNode = EachEdgeInfo.ToNode;
                long Cost = EachEdgeInfo.Cost;
                string EdgeHash = DeriveEdgeHash(FromNode, ToNode);
                mEdgeCostDict[EdgeHash] = Cost;
            }
        }

        // 最大コスト[状態]なDP表
        var TreeDPDict = new Dictionary<long, long>();

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

        foreach (JyoutaiDef1 EachJyoutai in DFSResult) {
            // 親ノードに配る木DP
            long Node = EachJyoutai.CurrNode;
            if (TreeDPDict.ContainsKey(Node) == false) {
                TreeDPDict[Node] = DeriveDefaultNodeVal();
            }

            long ParentNode = EachJyoutai.ParentNode;

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

            if (TreeDPDict.ContainsKey(ParentNode) == false) {
                TreeDPDict[ParentNode] = DeriveDefaultNodeVal();
            }
            long FromNode = EachJyoutai.CurrNode;
            long ToNode = EachJyoutai.ParentNode;
            TreeDPDict[ParentNode] =
                DeriveNewNodeVal(TreeDPDict[ParentNode],
                                 DeriveSendVal(TreeDPDict[Node], FromNode, ToNode));
        }

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

        // ReRootingなDFSを行う
        var sb = new System.Text.StringBuilder();
        foreach (long EachLong in ExecDFS2(1, TreeDPDict)) {
            sb.Append(EachLong);
            sb.AppendLine();
        }
        Console.Write(sb.ToString());
    }

    // 枝のハッシュ値を返す
    static string DeriveEdgeHash(long pFromNode, long pToNode)
    {
        long Min = Math.Min(pFromNode, pToNode);
        long Max = Math.Max(pFromNode, pToNode);
        return string.Format("{0},{1}", Min, Max);
    }

    // 枝のコスト[枝のハッシュ値]なDict
    static Dictionary<string, long> mEdgeCostDict = new Dictionary<string, 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 (mEdgeInfoListDict.ContainsKey(Popped.CurrNode) == false) {
                continue;
            }
            foreach (EdgeInfoDef EachToNode in mEdgeInfoListDict[Popped.CurrNode]) {
                if (VisitedSet.Add(EachToNode.ToNode)) {
                    WillPush.CurrNode = EachToNode.ToNode;
                    WillPush.ParentNode = Popped.CurrNode;
                    Stk.Push(WillPush);
                }
            }
        }
        return WillReturn;
    }

    // ReRootingなDFSを行う
    struct JyoutaiDef2
    {
        internal long PrevRootNode;
        internal long PrevRootVal;
        internal long CurrRootNode;
        internal long CurrRootVal;
    }
    static IEnumerable<long> ExecDFS2(long pRootNode, Dictionary<long, long> pTreeDPDict)
    {
        var WillReturn = new List<long>();

        var Stk = new Stack<JyoutaiDef2>();
        JyoutaiDef2 WillPush;
        WillPush.PrevRootNode = -1;
        WillPush.PrevRootVal = -1;
        WillPush.CurrRootNode = pRootNode;
        WillPush.CurrRootVal = pTreeDPDict[pRootNode];
        Stk.Push(WillPush);

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

        var AnswerDict = new Dictionary<long, long>();
        while (Stk.Count > 0) {
            JyoutaiDef2 Popped = Stk.Pop();

            AnswerDict[Popped.CurrRootNode] = Popped.CurrRootVal;

            // 累積値を管理する
            int ChildCnt = 0;
            if (mEdgeInfoListDict.ContainsKey(Popped.CurrRootNode)) {
                mEdgeInfoListDict[Popped.CurrRootNode].RemoveAll(pX => pX.ToNode == Popped.PrevRootNode);
                ChildCnt = mEdgeInfoListDict[Popped.CurrRootNode].Count;
            }
            if (ChildCnt == 0) continue;

            long[] RunValSei = new long[ChildCnt]; // 正方向の累積値
            long[] RunValRev = new long[ChildCnt]; // 逆方向の累積値

            long StaNode = mEdgeInfoListDict[Popped.CurrRootNode][0].ToNode;
            long EndNode = mEdgeInfoListDict[Popped.CurrRootNode][ChildCnt - 1].ToNode;
            RunValSei[0] =
                DeriveSendVal(pTreeDPDict[StaNode], StaNode, Popped.CurrRootNode);
            RunValRev[RunValRev.GetUpperBound(0)] =
                DeriveSendVal(pTreeDPDict[EndNode], EndNode, Popped.CurrRootNode);

            for (int I = 1; I <= ChildCnt - 1; I++) {
                long CurrChildNode = mEdgeInfoListDict[Popped.CurrRootNode][I].ToNode;
                RunValSei[I] =
                    DeriveNewNodeVal(RunValSei[I - 1],
                    DeriveSendVal(pTreeDPDict[CurrChildNode], CurrChildNode, Popped.CurrRootNode));
            }
            for (int I = ChildCnt - 2; 0 <= I; I--) {
                long CurrChildNode = mEdgeInfoListDict[Popped.CurrRootNode][I].ToNode;
                RunValRev[I] =
                    DeriveNewNodeVal(RunValRev[I + 1],
                    DeriveSendVal(pTreeDPDict[CurrChildNode], CurrChildNode, Popped.CurrRootNode));
            }

            // 状態の遷移
            for (int I = 0; I <= mEdgeInfoListDict[Popped.CurrRootNode].Count - 1; I++) {
                long CurrRootNode = mEdgeInfoListDict[Popped.CurrRootNode][I].ToNode;
                if (VisitedSet.Add(CurrRootNode) == false) {
                    continue;
                }
                if (CurrRootNode > mN) continue;

                // 左方向の累積値
                long LeftRunVal = DeriveDefaultNodeVal();
                bool HasLeftRunVal = false;
                if (I > 0) {
                    LeftRunVal = RunValSei[I - 1];
                    HasLeftRunVal = true;
                }

                // 右方向の累積値
                long RightRunVal = DeriveDefaultNodeVal();
                bool HasRightRunVal = false;
                if (I < mEdgeInfoListDict[Popped.CurrRootNode].Count - 1) {
                    RightRunVal = RunValRev[I + 1];
                    HasRightRunVal = true;
                }

                WillPush.PrevRootNode = Popped.CurrRootNode;
                var PrevRootValKouhoList = new List<long>();
                if (HasLeftRunVal) PrevRootValKouhoList.Add(LeftRunVal);
                if (HasRightRunVal) PrevRootValKouhoList.Add(RightRunVal);
                if (Popped.PrevRootNode > -1) {
                    PrevRootValKouhoList.Add(DeriveSendVal(
                        Popped.PrevRootVal, Popped.PrevRootNode, Popped.CurrRootNode));
                }
                WillPush.PrevRootVal = DeriveMergedSendVal(PrevRootValKouhoList);

                WillPush.CurrRootNode = CurrRootNode;
                var CurrRootValKouhoList = new List<long>();
                CurrRootValKouhoList.Add(pTreeDPDict[CurrRootNode]);
                CurrRootValKouhoList.Add(DeriveSendVal(
                    WillPush.PrevRootVal, WillPush.PrevRootNode, WillPush.CurrRootNode));
                WillPush.CurrRootVal = DeriveMergedSendVal(CurrRootValKouhoList);

                Stk.Push(WillPush);
            }
        }

        foreach (var EachPair in AnswerDict.OrderBy(pX => pX.Key)) {
            long Edge1Val = mDArr[EachPair.Key - 1];
            if (Edge1Val == EachPair.Value) {
                long Result = SolveNaive(EachPair.Key);
                yield return Result;
            }
            else {
                yield return EachPair.Value;
            }
        }
    }

    struct JyoutaiDefNaive
    {
        internal long CurrNode;
        internal long CostSum;
        internal long Level;
    }

    static long SolveNaive(long pRootNode)
    {
        mEdgeInfoListDict.Clear();

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

        Action<long, long, long> AddEdge = (pFrom, pTo, pCost) =>
        {
            if (mEdgeInfoListDict.ContainsKey(pFrom) == false) {
                mEdgeInfoListDict[pFrom] = new List<EdgeInfoDef>();
            }
            EdgeInfoDef WillAdd;
            WillAdd.ToNode = pTo;
            WillAdd.Cost = pCost;
            mEdgeInfoListDict[pFrom].Add(WillAdd);
        };

        foreach (string EachStr in mInputList.Skip(1).Take((int)mN - 1)) {
            SplitAct(EachStr);
            long A = wkArr[0];
            long B = wkArr[1];
            long C = wkArr[2];

            AddEdge(A, B, C);
            AddEdge(B, A, C);
        }

        var Stk = new Stack<JyoutaiDefNaive>();
        JyoutaiDefNaive WillPush;
        WillPush.CurrNode = pRootNode;
        WillPush.CostSum = 0;
        WillPush.Level = 0;
        Stk.Push(WillPush);

        var VistiedSet = new HashSet<long>();
        VistiedSet.Add(pRootNode);
        long MaxCost = long.MinValue;
        while (Stk.Count > 0) {
            JyoutaiDefNaive Popped = Stk.Pop();

            if (Popped.Level >= 1) {
                MaxCost = Math.Max(MaxCost, Popped.CostSum + mDArr[Popped.CurrNode - 1]);
            }

            foreach (EdgeInfoDef EachEdgeInfo in mEdgeInfoListDict[Popped.CurrNode]) {
                if (VistiedSet.Add(EachEdgeInfo.ToNode)) {
                    WillPush.CurrNode = EachEdgeInfo.ToNode;
                    WillPush.CostSum = Popped.CostSum + EachEdgeInfo.Cost;
                    WillPush.Level = Popped.Level + 1;
                    Stk.Push(WillPush);
                }
            }
        }
        return MaxCost;
    }
}

解説

下記の手順で解いてます。

手順01
ノードと枝の両方がコストを持つと
厄介なので、葉ノードを用意し、
枝のみがコストを持つようする

手順02
木DPとReRootingで根ノードごとの
最大コストを求める

手順03
根ノードのコストが、枝が1本だけのコストの場合は、
ナイーブな方法で、正しいコストを求める