[AtCoder - ABC243 F] Lottery

난이도: 2125

 

태그

• Combinatorics
  • Addition Rule
• Probability Theory
  • Dynamic Programming using Probability
• Modular Multiplicative Inverse

 

풀이

1. 앞에 있는 \( N \)개만 두고 볼 때, \( K \)번 추첨해서 \( M \)개를 뽑을 확률은?

\( dp_{n, m, k} \) = 앞에 있는 \( n \)개만 두고 볼 때, \( k \)번 추첨해서 \( m \)개를 뽑을 확률로 정의해봅시다.

 

그리고, 이에 맞게 초항부터 설정해봅시다.

• \( n < m \)인 경우: 물건 개수보다 더 많은 걸 뽑았다는데, 당연히 불가능합니다.
• \( m > k \)인 경우: 추첨 횟수보다 더 많이 뽑았다는데, 불가능합니다.
• \( m = k = 0 \)인 경우: 추첨을 안 했으니, 당연히 아무것도 못 뽑았겠죠. 이 때의 확률은 1입니다.
• \( m = 0; k \neq 0 \)인 경우: 추첨을 했는데 아무것도 못 뽑았다고 하는데, 불가능합니다.
• \( n = 1 \)인 경우: \( m, k = 0 \)이거나 \( m, k \neq 0 \)이면 1, 아니면 0입니다.

 

그리고, 점화식을 설정해봅시다.

\( k \)번의 추첨 과정에서 \( n \)번째 물건을 \( c \)번 뽑았다고 할 때, 아래와 같이 나눠볼 수 있습니다.

• \( c = 0 \): \( n \)번째 물건을 뽑지 못했습니다.
  이는 \( dp_{n-1, m, k} \)로 나타낼 수 있겠죠.
  하지만, \( dp_{n-1, m, k} \)는 \( n-1 \)개의 물건으로 추첨하는 확률을 가지고 있기 때문에
  이를 \( n \)개의 물건으로 추첨하는 확률로 변경시키기 위해
  분모값인 \( \left( \sum\limits_{i=1}^{n-1} W_{i} \right)^k \)를 \( \left( \sum\limits_{i=1}^{n} W_{i} \right)^k \)로 바꿔줘야 합니다.
  // 이는 이전 분모 값에 해당하는 값을 곱한 뒤, 현재 분모 값에 해당하는 값으로 나눠서 해결할 수 있습니다.

 

• \( c \neq 0 \): \( n \)번째 물건을 1회 이상 뽑은 경우입니다.
  \( n \)번째 물건을 뽑은 추첨을 제외한다면, 남은 경우는 \( dp_{n-1, m-1, k-c} \)로 나타낼 수 있습니다.
  물론, 위에서 했던 처리와 동일하게 \( n \)개의 물건으로 추첨하는 확률로 바꿔야겠죠.
  // 이 때에는 지수가 \( k-c \)로 들어가야 합니다!
  그 뒤로는 \( n \)번째 물건을 \( c \)번 뽑을 확률인 \( \left( \frac{W_{n}}{\sum\limits_{i=1}^{n} W_{i}} \right)^c \)를 곱해주고,
  \( n \)번째 물건을 뽑았을 때와 아닐 때를 적절히 섞어주는 경우의 수인 \( _kC_c \)를 곱해주면 됩니다.

 

2. 코드

const ll mod = 998244353;
ll fpow(ll mul, ll bit){ ll res = 1;
    while (bit){
        if (bit&1){ res = res*mul % mod; }
        mul = mul*mul % mod; bit >>= 1;
    }
    return res;
}
inline ll finv(ll x){ return fpow(x, mod-2); }
 
const int N = 50;
ll fac[52], inv[52];
inline ll ncr(int n, int r){
    if (0 > r || r > n){ return 0; }
    return fac[n] *inv[r] %mod *inv[n-r] %mod;
}
 
ll arr[52], prf[52];
 
ll dp[52][52][52];
ll dpf(int n, int m, int k){
    if (dp[n][m][k] != -1){ return dp[n][m][k]; }
    if (n < m || m > k){ return dp[n][m][k] = 0; }
    if (m == 0){ return dp[n][m][k] = (k==0); }
    if (n == 1){ return dp[n][m][k] = ((m==0) == (k==0)); }
    ll val = prf[n-1]*finv(prf[n]) % mod;
    dp[n][m][k] = dpf(n-1, m, k)*fpow(val, k) % mod;
    for (int c = 1; c <= k; c++){
        ll res = dpf(n-1, m-1, k-c)*fpow(val, k-c) % mod;
        res = res*fpow(arr[n]*finv(prf[n]) % mod, c) % mod;
        res = res*ncr(k, c) % mod;
        dp[n][m][k] += res; dp[n][m][k] %= mod;
    }
    //cout << "DP " << n << ' ' << m << ' ' << k << " = " << dp[n][m][k] << endl;
    return dp[n][m][k];
}
 
void Main(){
    fac[0] = 1; for (int i = 1; i <= N; i++){ fac[i] = fac[i-1]*i % mod; }
    inv[N] = finv(fac[N]); for (int i = N; i >= 1; i--){ inv[i-1] = inv[i]*i % mod; }
    int n, m, k; cin >> n >> m >> k;
    for (int i = 1; i <= n; i++){ cin >> arr[i]; }
    for (int i = 1; i <= n; i++){ prf[i] = prf[i-1] + arr[i]; }
    memset(dp, -1, sizeof(dp));
    cout << dpf(n, m, k);
}