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でノードごとに、(自身も含めた)子孫ノードの数を求めます。
それから、枝を全探索し、
子側のノードの合計と親側のノードの合計で、積の法則を使い、解に計上します。