CS Academy BST Fixed Height

リンク : https://csacademy.com/contest/archive/task/bst-fixed-height/statement/

問題概要

Nこの頂点を持ち、深さHの完全二分木の一部を使った木は何通りあるか求めよ。

解説

正直自分にはこの解説はかけない。最初、dp[i][j][k]:深さiのときの頂点をj個をすでに使い、i-1の頂点がk個あるときと考えたが、到底間に合わない。そこで求めるべきdpは[i][j]でi個すでに使い、深さjのときの通り数である。そこで、dp[i][j]の通り数は右の部分木と左の部分木のどちらかがj-1であればよい。両方j-1のときと、片方がj-1でありもう片方はj-1未満であるときを求めればよい。
両方j-1のときは
dp[i][j]+=sum{k=0,1..i}{dp[k][j-1]×dp[i-k][j-1]}
片方j-1のときは
dp[i][j]+=sum{k=0,1...i}{dp[k][j-1]×sum{l=0,1...j-2}{dp[i-k][l]}}
ここで計算量を減らすために、sum{l=0,1...j-2}{dp[i-k][l]}の部分は別で値を持っておく必要がある。これでO(H*N²)で解けた。
しかし実は、このまま実装すると数ケースでTLEするので自明に必要のないforの範囲を縮めたりする必要がある。

LL n, h, dp1[600][600], dp2[600][600];
LL DP(int x, int y, int flag) {
    if (flag == 1) {
        if (y < 0)return 0;
        if (x < y)return 0;
        if (y == 1)return x == 1;
        if (dp1[x][y] != -1)return dp1[x][y];
        LL ret = 0;
        for (int i = y - 1; x - i - 1 >= y - 1; i++) {
            ret += DP(i, y - 1, 1)*DP(x - i - 1, y - 1, 1);
            ret %= MOD;
        }
        for (int i = y - 1; i < x; i++) {
            ret += DP(i, y - 1, 1)*DP(x - i - 1, y - 2, 2) * 2;
            ret %= MOD;
        }
        return dp1[x][y] = ret;
    }
    else {
        if (dp2[x][y] != -1)return dp2[x][y];
        if (x == 0)return 1;
        if (x < 0||y < 0)return 0;
        LL ret = 0;
        ret += DP(x, y, 1) + DP(x, y - 1, 2);
        ret %= MOD;
        return dp2[x][y] = ret;
    }
}
int main() {
    cin >> n >> h;
    rep(i, n + 1)rep(j, h + 1)dp1[i][j] = -1;
    rep(i, n + 1)rep(j, h + 1)dp2[i][j] = -1;
    cout << DP(n, h, 1) << endl;
    return 0;
}