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

問題

方針

一見難しそうに見えますが、\( 300 \) 点の問題なので解法は簡単であることが予想できます。

実現不可能な整数列

  • 頂点 \( 1 \) からの距離について

頂点 \( 1 \) からの距離が与えられるので、\( D_1 = 0 \) が必要です。

  • \( N – 1 \) 辺が存在する

\( N – 1 \) 辺存在することから、\( D_i \neq 0 \ ( i \neq 1)\) が必要です。つまり、距離が \( 0 \) となる頂点は頂点 \( 1 \) だけです。

  • 距離の連続性

距離が \( i \) である頂点の個数を \( d_i \) とします。\( d_i \neq 0 \) を満たす最大の \( i \) を \( j \) とします。このとき、\( d_i \neq 0 \ (1 \leq i \leq j – 1) \) を満たす必要があります。例えば、距離が \( 9 \) の頂点が存在するとき、\( 8 \) 以下の距離について、少なくとも一つ以上頂点が存在する必要があります。

距離の頻度に着目

距離が \( i \) である頂点の個数を \( d_i \) とします。また、\( D_1 = 0 \) であり、\( d_0 = 1 \) であるとします。距離 \( 1 \) から順に考えていきます。

距離 \( 1 \) までの木を考えると、接続される頂点は \( 1 \) なので、\( d_1 \) の値に依らず木は一意に定まります。次に、距離 \( 2 \) までの木を考えると、接続される頂点の候補は \( d_1 \) 個あるので、\( d_1^{d_2} \) 個の木が考えられます。同様にして距離 \( N – 1 \) まで考えると、求める木の個数は、

\[d_0^{d_1} \times d_1^{d_2} \times \cdots \times d_{N-2}^{d_{N-1}}\]

となります。また、計算過程で階乗の剰余を行う必要があります。

コード

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;

// a^n mod を計算する
ll modpow(ll a, ll n, ll mod) {
    ll res = 1;
    while (n > 0) {
        if (n & 1) res = res * a % mod;
        a = a * a % mod;
        n >>= 1;
    }
    return res;
}

int main() {
    int N;
    cin >> N;
    int D[N];
    ll mod = 998244353;
    ll d[N]{};
    for (int i = 0; i < N; i++) {
        cin >> D[i];
        d[D[i]]++;
    }
    ll ans = 1;
    bool flag = false;
    for (int i = N - 1; i >= 1; i--) {
        if (d[i] != 0) {
            flag = true;
        }
        if (flag && d[i] == 0) {
            cout << "0\n";
            return 0;
        }
    }
    if (d[0] != 1) {
        cout << "0\n";
        return 0;
    }
    if (D[0] != 0) {
        cout << "0\n";
        return 0;
    }
    for (int i = 1; i < N; i++) {
        ans *= modpow(d[i - 1], d[i], mod);
        ans %= mod;
    }
    cout << ans << "\n";
    return 0;
}