mayoko’s diary

プロコンとかいろいろ。

POJ 3046 Ant Counting

この解法 TLE すると思ってたので「は?」と言わざるを得ない。

解法

dp[n][t] = (n 種類目まで見ている時点で, セットの大きさが t 以下になるような集合の数)とします。
すると, dp[n][t] は, n 種類目まででセットの大きさが t-1 以下になっているか, t になっているかのいずれかです。また, セットの大きさが t になるのは, n-1 種類目まででのセットの大きさが p であるような場合の数を way[p] として,
way[t-cnt[n]] + way[t-cnt[n]+1] + ... + way[t]
で表せます(cnt[n] は種類が n であるような蟻の数)。

これは, 上で定義した dp では dp[n-1][t] - dp[n-1][t-cnt[n]-1] で表せるので, 高速に計算できます。

ただ, それでも計算量は O(A*T) で, A は10^3, T は10^5 までありえるので, ジャッジが遅い POJ だと TLE すると思ったのですが。まぁ通れば正義ということにします。

const int MAXT = 1024;
const int MOD = 1000000;
int cnt[MAXT];
int dp[2][MAXT*100];

int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    int T, A, S, B;
    scanf("%d%d%d%d", &T, &A, &S, &B);
    for (int i = 0; i < A; i++) {
        int tmp;
        scanf("%d", &tmp);
        cnt[tmp-1]++;
    }
    for (int i = 0; i <= A; i++) dp[0][i] = 1;
    for (int i = 0; i < T; i++) {
        int cur = i%2;
        int tar = 1^cur;
        dp[tar][0] = 1;
        for (int j = 1; j <= A; j++) {
            dp[tar][j] = dp[tar][j-1];
            dp[tar][j] += dp[cur][j];
            if (j-cnt[i]-1 >= 0) dp[tar][j] -= dp[cur][j-cnt[i]-1];
            dp[tar][j] %= MOD;
            if (dp[tar][j] < 0) dp[tar][j] += MOD;
        }
    }
    int ans = dp[T%2][B];
    if (S > 0) ans -= dp[T%2][S-1];
    if (ans < MOD) ans += MOD;
    cout << ans%MOD << endl;
    return 0;
}