yukicoder No.243 出席番号(2)

これは良い問題だと思いました。

問題

No.243 出席番号(2) - yukicoder

解法

なんとなく包除原理っぽいです。

(0個以上嫌いナンバーが選ばれる場合の数) - (1個以上嫌いナンバーが選ばれる場合の数) + (2個以上嫌いナンバーが選ばれる場合の数) - ...

という感じで答えを求めることを考えます。

0個以上嫌いナンバーが選ばれる場合の数は,で求められます。また,1個以上嫌いナンバーが選ばれる場合の数は,
(0が嫌いな人の数) * +
(1が嫌いな人の数) * + ...

2個以上嫌いナンバーが選ばれる場合の数は,
(0が嫌いな人の数) * (1が嫌いな人の数) * + ...

というような感じです。

包除原理で考える上で大事なのは「嫌いナンバーとして選ばれる数字の個数」です。これを基準にdp[k] := (嫌いナンバーとして選ばれる数字の個数がk個の場合の数)を求めると,上記のような包除原理を適用でき,答えを求めることが出来ます。

const int MAXN = 5005;
const ll MOD = 1e9+7;
int N;
int cnt[MAXN];

ll dp[2][MAXN];
ll fact[MAXN];

int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);
    fact[0] = 1;
    for (int i = 1; i < MAXN; i++) {
        fact[i] = (fact[i-1] * i) % MOD;
    }
    cin >> N;
    for (int i = 0; i < N; i++) {
        int a;
        cin >> a;
        cnt[a]++;
    }
    dp[0][0] = 1;
    for (int i = 0; i < N; i++) {
        int cur = i%2;
        int tar = (1^cur);
        memset(dp[tar], 0, sizeof(dp[tar]));
        for (int j = 0; j <= i; j++) {
            (dp[tar][j] += dp[cur][j]) %= MOD;
            (dp[tar][j+1] += dp[cur][j] * cnt[i]) %= MOD;
        }
    }
    ll ans = 0;
    for (int i = 0; i <= N; i++) {
        ll plus = (dp[N%2][i]*fact[N-i]) % MOD;
        if (i%2) plus *= -1;
        (ans += plus) %= MOD;
    }
    if (ans < 0) ans += MOD;
    cout << ans << endl;
    return 0;
}