No.1035 Color Box

yukicoder.me

$ M $ 種類のペンキをすべて使わなくても良い場合は、 $ N $ 個の箱がそれぞれが $ M $ 個の選択肢を持っているため $M ^{N}$ 通りとなります。

ここで問題となってしまうのが、$ M $ 種類のペンキのうち1回も使われないものが存在することです。

そこで、

$$ (0種類以上のペンキを使わない) - (1種類以上のペンキを使わない) + (2種類以上のペンキを使わない) $$

といった具合で包除原理を適用します。

$ i $ 種類以上のペンキを使わない場合の数は、

  • $ M $ 種類のペンキのうち使用する $ M - i $ 種類を選ぶ $ \,_{M} C_{M - i} $
  • $ N $ 個の箱の色を $ M - i $ 種類のペンキのうちから選ぶ $ (M - i)^{N} $

の積をとることで、

$$ _{M}C_{M - i} \times (M - i) ^ {N} $$

となります。したがって、全体の場合の数は

$$ \sum _{i = 0}^{M - 1} (-1)^{i} \times _{M}C_{M - i} \times (M - i) ^ {N} $$

と求めることが出来ます。 $ M - i $種類以上のペンキを使用しないとすると、

$$ \sum _{i = 1}^{M} (-1)^{M - i} \times _{M}C_{i} \times i ^ {N} $$

とすることもできます。

#include <atcoder/all>
using namespace std;
using mint = atcoder::modint1000000007;

int main(){
    int n, m;
    cin >> n >> m;
    
    vector<mint> fact(m + 1), inv(m + 1);
    fact[0] = 1;
    for(int i = 1; i <= m; i++) fact[i] = i * fact[i - 1];
    inv[m] = 1 / fact[m];
    for(int i = m; i >= 1; i--) inv[i - 1] = i * inv[i];
    
    mint ans;
    for(int i = 1; i <= m; i++){
        ans += (m - i & 1 ? -1 : 1) * inv[m - i] * inv[i] * mint(i).pow(n);
    }
    
    cout << (ans * fact[m]).val() << '\n';
}