第二回全国統一プログラミング王決定戦予選 B Counting of Trees

問題へのリンク

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("0 1 1 2");
            //2
        }
        else if (InputPattern == "Input2") {
            WillReturn.Add("4");
            WillReturn.Add("1 1 1 1");
            //0
        }
        else if (InputPattern == "Input3") {
            WillReturn.Add("7");
            WillReturn.Add("0 3 2 1 2 2 1");
            //24
        }
        else {
            string wkStr;
            while ((wkStr = Console.ReadLine()) != null) WillReturn.Add(wkStr);
        }
        return WillReturn;
    }

    const long Hou = 998244353;

    static void Main()
    {
        List<string> InputList = GetInputList();
        int[] DArr = InputList[1].Split(' ').Select(pX => int.Parse(pX)).ToArray();

        // 頂点1の距離が0でない場合は不可
        if (DArr[0] != 0) {
            Console.WriteLine(0);
            return;
        }

        // 頂点1以外に距離が0の頂点があったら不可
        for (int I = 1; I <= DArr.GetUpperBound(0); I++) {
            if (DArr[I] == 0) {
                Console.WriteLine(0);
                return;
            }
        }

        Dictionary<int, int> CntDict =
            DArr.GroupBy(pX => pX).ToDictionary(pX => pX.Key, px => px.Count());

        long Answer = 1;
        foreach (var EachPair in CntDict) {
            if (EachPair.Key == 0) continue;

            int PrevKey = EachPair.Key - 1;
            if (CntDict.ContainsKey(PrevKey)) {
                long Bekijyou = 1;
                for (int I = 1; I <= EachPair.Value; I++) {
                    Bekijyou *= CntDict[PrevKey];
                    Bekijyou %= Hou;
                }
                Answer *= Bekijyou;
                Answer %= Hou;
            }
            else {
                Answer = 0;
                break;
            }
        }
        Console.WriteLine(Answer);
    }
}

解説

0 1 1 2
というサンプルで考察すると
1の親は1通り
2の親は2通り
で、積の法則の1*2で
木は2通りだと分かります。

0 3 2 1 2 2 1
というサンプルで考察すると
1の親は1通り
2の親は2*2*2通り で 8通り
3の親は3通り
で、積の法則の8*3で
木は24通りだと分かります。