ABC479-E Cut and Add

問題へのリンク

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("2 4");
            WillReturn.Add("1 5");
            WillReturn.Add("1 4");
            WillReturn.Add("1 3");
            //10 6 6 7 8
        }
        else if (InputPattern == "Input2") {
            WillReturn.Add("3");
            WillReturn.Add("1 2");
            WillReturn.Add("2 3");
            //2 3 1
        }
        else if (InputPattern == "Input3") {
            WillReturn.Add("9");
            WillReturn.Add("4 6");
            WillReturn.Add("2 3");
            WillReturn.Add("1 5");
            WillReturn.Add("7 8");
            WillReturn.Add("2 9");
            WillReturn.Add("5 6");
            WillReturn.Add("2 6");
            WillReturn.Add("7 9");
            //20 36 34 28 23 29 23 19 31
        }
        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 AllNodeCnt = long.Parse(InputList[0]);

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

        foreach (string EachStr in InputList.Skip(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);

        // DFS順での双対セグ木
        var InsDualSegmentTree = new DualSegmentTree(AllNodeCnt);

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

            // 枝の両端ノードの、親子関係を判定
            ClassTreeAnalyze.NodeInfoDef NodeInfo1 = mCurrNodeDict[FromNode];
            ClassTreeAnalyze.NodeInfoDef NodeInfo2 = mCurrNodeDict[ToNode];

            long ParentNode, ChildNode;
            if (NodeInfo1.CurrNode == NodeInfo2.ParentNode) {
                ParentNode = NodeInfo1.CurrNode;
                ChildNode = NodeInfo2.CurrNode;
            }
            else {
                ParentNode = NodeInfo2.CurrNode;
                ChildNode = NodeInfo1.CurrNode;
            }

            ClassTreeAnalyze.NodeInfoDef ParentNodeInfo = mCurrNodeDict[ParentNode];
            ClassTreeAnalyze.NodeInfoDef ChildNodeInfo = mCurrNodeDict[ChildNode];

            long ChildNodeCnt = ChildNodeInfo.DFSNoEnd - ChildNodeInfo.DFSNoSta + 1;
            long OtherNodeCnt = AllNodeCnt - ChildNodeCnt;

            long RangeSta = ChildNodeInfo.DFSNoSta;
            long RangeEnd = ChildNodeInfo.DFSNoEnd;

            if (ChildNodeCnt > OtherNodeCnt) {
                InsDualSegmentTree.RangeAdd(RangeSta, RangeEnd, AddVal);
            }
            else {
                InsDualSegmentTree.RangeAdd(0, AllNodeCnt, AddVal);
                InsDualSegmentTree.RangeAdd(RangeSta, RangeEnd, -AddVal);
            }
            AddVal++;
        }

        var AnswerList = new List<long>();
        for (long I = 1; I <= AllNodeCnt; I++) {
            ClassTreeAnalyze.NodeInfoDef CurrNodeInfo = mCurrNodeDict[I];
            long SegInd = CurrNodeInfo.DFSNoCurr;
            AnswerList.Add(InsDualSegmentTree[SegInd]);
        }
        Console.WriteLine(LongEnumJoin(" ", AnswerList));
    }

    // セパレータとLong型の列挙を引数として、結合したstringを返す
    static string LongEnumJoin(string pSeparater, IEnumerable<long> pEnum)
    {
        string[] StrArr = Array.ConvertAll(pEnum.ToArray(), pX => pX.ToString());
        return string.Join(pSeparater, StrArr);
    }
}

#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での訪問順]なDict
    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<JyoutaiDefTreeAnalyze>();
        JyoutaiDefTreeAnalyze 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) {
            JyoutaiDefTreeAnalyze 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(pNodeInfoList);
    }

    private struct JyoutaiDefTreeAnalyze
    {
        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

// 区間加算、1点取得な双対セグ木(フェニック木使用)
#region DualSegmentTree
internal class DualSegmentTree
{
    private long[] mBitArr; // 内部配列(1オリジンなため、添字0は未使用)
    private long mExternalArrUB;

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

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

    // コンストラクタ
    // フェニック木の外部配列(0オリジン)のUBを指定
    internal DualSegmentTree(long pExternalArrUB)
    {
        mExternalArrUB = pExternalArrUB;

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

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

    // 双対セグメント木の機能
    // 区間加算
    internal void RangeAdd(long pSta, long pEnd, long AddVal)
    {
        pSta++; // 1オリジンに変更
        pEnd++; // 1オリジンに変更

        long ImosSta = pSta;
        long ImosEnd = pEnd + 1;

        // いもす法
        FenwickTree_Add(ImosSta, AddVal);
        if (ImosEnd <= mBitArr.GetUpperBound(0)) {
            FenwickTree_Add(ImosEnd, -AddVal);
        }
    }

    // 双対セグメント木の機能
    // 1点取得
    internal long GetVal(long pInd)
    {
        pInd++; // 1オリジンに変更
        return FenwickTree_GetSum(1, pInd);
    }

    // フェニック木の機能
    // [pSta,pEnd] のSumを返す
    private long FenwickTree_GetSum(long pSta, long pEnd)
    {
        return FenwickTree_GetSum(pEnd) - FenwickTree_GetSum(pSta - 1);
    }

    // フェニック木の機能
    // [0,pEnd] のSumを返す
    private long FenwickTree_GetSum(long pEnd)
    {
        long Sum = 0;
        while (pEnd >= 1) {
            Sum += mBitArr[pEnd];
            pEnd -= pEnd & -pEnd;
        }
        return Sum;
    }

    // フェニック木の機能
    // [I] に Xを加算
    private void FenwickTree_Add(long pI, long pX)
    {
        while (pI <= mBitArr.GetUpperBound(0)) {
            mBitArr[pI] += pX;
            pI += pI & -pI;
        }
    }
}
#endregion

解説

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

それから、区間加算、1点取得な双対セグ木を作り、
枝ごとに、子ノードのノード数と、子ノード以外のノード数を求め、
双対セグ木の区間加算でシュミレーションしてます。