E8本(数学)    前のE8本(数学)の問題へ

E8本(数学) 099 Tree Distance(★5)


問題へのリンク


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

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

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

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

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

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

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

        List<JyoutaiDef> DFSResult = ExecDFS();
        DFSResult.Reverse();

        // 部分木のノード数[ノード]なDict
        var NodeCntDict = new Dictionary<int, int>();

        for (int I = 1; I <= N; I++) {
            NodeCntDict[I] = 1;
        }

        // 葉から根へ、配る木DP
        foreach (JyoutaiDef EachJyoutai in DFSResult) {
            if (EachJyoutai.ParentNode == -1) {
                continue;
            }
            NodeCntDict[EachJyoutai.ParentNode] += NodeCntDict[EachJyoutai.Node];
        }

        // 葉から根への、辺の全探索
        long Answer = 0;
        foreach (JyoutaiDef EachJyoutai in DFSResult) {
            if (EachJyoutai.ParentNode == -1) {
                continue;
            }

            long NodeCnt1 = NodeCntDict[EachJyoutai.Node];
            long NodeCnt2 = N - NodeCnt1;

            // 積の法則
            Answer += NodeCnt1 * NodeCnt2;
        }
        Console.WriteLine(Answer);
    }

    struct JyoutaiDef
    {
        internal int Node;
        internal int ParentNode;
    }

    static List<JyoutaiDef> ExecDFS()
    {
        var WillReturn = new List<JyoutaiDef>();

        var Stk = new Stack<JyoutaiDef>();
        JyoutaiDef WillPush;
        WillPush.Node = 1;
        WillPush.ParentNode = -1;
        Stk.Push(WillPush);

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

        while (Stk.Count > 0) {
            JyoutaiDef Popped = Stk.Pop();
            WillReturn.Add(Popped);
            if (mToNodeListDict.ContainsKey(Popped.Node)) {
                foreach (int EachToNode in mToNodeListDict[Popped.Node]) {
                    if (VisitedSet.Add(EachToNode)) {
                        WillPush.Node = EachToNode;
                        WillPush.ParentNode = Popped.Node;
                        Stk.Push(WillPush);
                    }
                }
            }
        }
        return WillReturn;
    }
}


解説

最初に、木DPでノードごとに、(自身も含めた)子孫ノードの数を求めます。
それから、枝を全探索し、
子側のノードの合計と親側のノードの合計で、積の法則を使い、解に計上します。