【今日の精進】ABC215G - Colorful Candies 2

毎日格上の問題を倒すやつの16日目です。

問題リンク

atcoder.jp

問題概要

個の飴があるよ。

「この中から 個選んだ時の種類数の期待値」をすべての に対して求めてね。

解法

主客転倒テクです。 個ある飴を選ぶ確率を全部に関して足せば期待値になります。¹

こういうのはいつもの二項係数の組み合わせで求めることができます。

また、個数の種類数²が実はたかだか 個しかないことが証明できるため、 時間で全部計算できます。

提出コード

#include <bits/stdc++.h>
#include <atcoder/all>
using namespace std;
using namespace atcoder;
using mint = modint998244353;
#pragma GCC target("avx2")
#pragma GCC optimize("O3")
#pragma GCC optimize("unroll-loops")

vector<mint> fact(100000,1);
vector<mint> rev(100000,1);
mint comb(int n,int i,int r){
    if(n-i<r) return 0;
    if(n-i<0||r<0) return 0;
    return fact[n-i]*rev[n-i-r]*rev[n]*fact[n-r];
}

int main(){
    int N;
    cin>>N;
    for(int i=1; i<100000; i++){
        fact[i]=fact[i-1]*i;
        rev[i]=1/fact[i];
    }
    map<int,int> c;
    for(int i=0; i<N; i++){
        int c_i;
        cin>>c_i;
        c[c_i]++;
    }
    map<int,int> cnt;
    for(auto const& [_,i]:c){
        cnt[i]++;
    }
    vector<mint> ans(N,0);
    for(auto const& [i,j]:cnt){
        for(int k=0; k<N; k++){
            ans[k]+=j*(1-comb(N,i,k+1));
        }
    }
    for(mint i:ans){
        cout<<i.val()<<endl;
    }
}

感想

TLE除くのにめっちゃ時間かかった(激怒)³

期待値の線形性はもはや実家より安心する

---

1. #いつもの↩
2. 我ながらKAIBUNsyo↩
3. Python捨てても苦戦した↩