Intel Code Challenge Final Round (Div. 1 + Div. 2, Combined) F - Uniformly Branched Trees 无根树->有根树+dp

F - Uniformly Branched Trees

#include<bits/stdc++.h>
#define LL long long
#define fi first
#define se second
#define mk make_pair
#define PII pair<int, int>
#define PLI pair<LL, int>
#define ull unsigned long long
using namespace std;

const int N = 1e3 + 7;
const int inf = 0x3f3f3f3f;
const LL INF = 0x3f3f3f3f3f3f3f3f;

LL dp[N][11][N], inv[N], f[N], finv[N], g[N][11];
int n, d, mod;

void init() {
    inv[1] = f[0] = finv[0] = 1;
    for(int i = 2; i < N; i++) inv[i] = (mod-mod/i)*inv[mod%i]%mod;
    for(int i = 1; i < N; i++) f[i] = f[i-1]*i%mod;
    for(int i = 1; i < N; i++) finv[i] = finv[i-1]*inv[i]%mod;
}

int main() {
    scanf("%d%d%d", &n, &d, &mod);
    init();
    for(int i = 0; i <= n; i++) dp[1][0][i] = 1;
    for(int i = 1; i <= d; i++) g[1][i] = 1;
    for(int i = 2; i <= n; i++) {
        for(int j = 1; j <= d; j++) {
            for(int k = 1; k < i; k++) {
                for(int l = 1; l <= j && l*k <= i; l++) {
                    dp[i][j][k] = (dp[i][j][k] + 1ll*dp[i-l*k][j-l][k-1]*g[k][l])%mod;
                }
            }
            for(int k = 1; k <= n; k++) dp[i][j][k] = (dp[i][j][k]+dp[i][j][k-1])%mod;
        }
        g[i][1] = dp[i][d-1][n];
        for(int j = 2; j <= d; j++)
            g[i][j] = g[i][j-1]*(dp[i][d-1][n]+j-1)%mod*inv[j]%mod;
    }
    LL ans = 0;
    if(n <= 2) ans = 1;
    else ans = dp[n][d][n/2];
    if(n > 2 && !(n&1)) {
        ans = (ans + mod - dp[n/2][d-1][n/2] * (dp[n/2][d-1][n/2]-1) / 2 % mod) % mod;
    }
    printf("%lld\n", ans);
    return 0;
}

/*
*/