Codeforces Edu78 F [Cards] 풀이

문제

지문은 생략하고, 다음을 계산하는 문제이다.

Binomial Distribution $B(n, p)$ 가 주어진다. 이때, $E(X^k)$ 를 빠르게 구하고 싶다.

 

Dynamic Programming 을 이용한 풀이가 있는것 같은데, 사실 잘 모르겠다. 그 풀이를 설명한 댓글에서도 Stirling number of the second kind 같은 것들을 언급하는 것으로 보아 그닥 프로그래밍적인 풀이인지는 모르겠고, 내가 모르는 조합 / 확률론적인 지식을 요구하는 것 같아서 굳이 그걸 찾아서 공부해야 할지는...

 

아무튼 내 풀이는 moment generating function 과 미적분 지식을 요구하기 때문에 그닥 Codeforce에 적합한 것 같지는 않다. 다른 풀이의 존재를 몰랐다면 이거 그냥 구데기 문제라고 생각했을 듯..

 

풀이

Higher moment $E(X^k)$ 를 구해야 할 때 가장 먼저 생각해 볼 수 있는 것 중 하나는 적률생성함수 (Moment generating function, 이하 MGF) 이다. 확률변수 X의 MGF는 다음과 같이 정의된다.

$$M_X(t) = E(e^{tX})$$

이렇게 이상하게 정의된 함수가 쓸모가 있는 이유는, $e^x$ 의 테일러 전개식

$$e^x = 1 + x + \frac{x^2}{2!} + \dots$$ 때문이다. 이 식으로부터, $M_X^{(k)} (0) = E(X^k)$ 를 얻는다.

또한, Binomial Distribution의 Moment Generating Function이 다음과 같음이 잘 알려져 있다.

$$M(t) = (pe^t + 1 - p)^n$$

이제, 결국 위 식을 빠르게 잘 Evaluate 하는 방법을 찾으면 된다. 즉, 저 함수 $M(t)$ 를 k번 미분한 다음 $t = 0$ 을 대입, 즉 $e^t = 1$ 을 대입하면 된다.

 

$pe^t + 1 - p)^n$ 를 $f_0(t)$ 라고 하자 $f_0(t)$ 를 미분하면 $pne^t (pe^t + 1-p)^{n-1}$ 를 얻는데, 이를 $f_1 (t)$ 라고 하자. 이를 다시 미분하면 자기 자신과 함께 $p^2 n(n-1) e^{2t} (pe^t + 1-p)^{n-2}$ 를 얻는다. 

구체적으로 $f_i (t)$ 를 다음과 같이 정의한다.

$$f_i (t) = p^i n(n-1) \dots (n-i+1) e^{it} (pe^t + 1 - p) ^{n-i}$$ 

이제 이렇게 놓고 미분해 보면 $f_i'(t) = if_i(t) + f_{i+1} (t)$ 를 얻는다. 처음에 $f_0(t)$ 를 하나 들고 시작했으므로, $k$ 번의 iteration 이후에 남는 결과는 $p$와 $n$에 대한 $k$ 차 다항식이다. 각각의 계수는 이전 회차의 다항식을 한번 돌면서 계산할 수 있으므로, 이 과정의 시간 복잡도는 $O(k^2)$ 이다. 

 

결과 식에 남은 $f_i(0)$ 은 $p^i n(n-1)(n-2) \dots (n-i+1)$ 이 되므로, 이를 유리수를 이용하여 (각각의 수를 문제에 나온 대로 modulo 998244353을 이용하여 분수 해싱을 하는 식으로 할 수 있다) 정확히 evaluate 하면 된다.

 

#include <bits/stdc++.h>
using namespace std;
#define ll long long
#define usecppio ios::sync_with_stdio(0);cin.tie(0);cout.tie(0);

#define int ll
const int MOD = 998244353;
ll modpow(ll x, ll y, ll p)
{
	ll res = 1;
	x = x % p;
	while (y > 0)
	{
		if (y & 1)
			res = (res*x) % p;
		y = y>>1;
		x = (x*x) % p;
	}
	return res;
}

ll modinv(ll x, ll p)
{
	return modpow(x,p-2,p);
}

vector <int> v;
vector <int> v2;

void iter()
{
	for (int i = 1; i<=5000; i++)
	{
		v2[i+1] += v[i];
		v2[i+1] %= MOD;
		v2[i] += (v[i]*i);
		v2[i] %= MOD;
	}
	for (int i = 1; i<=5000; i++)
	{
		v[i] = v2[i];
		v2[i] = 0;
	}
}

int n, m, k;
int eval(int t, int p)
{
	if (p > n) return 0;
	int res = 1;
	for (int i = n; i>=(n-p+1); i--)
	{
		res *= i;
		res %= MOD;
	}
	int denom = modinv(modpow(m,p,MOD),MOD);
	res *= denom; res %= MOD;
	res *= t; res %= MOD;
	return res;
}
int32_t main()
{
	usecppio
	cin >> n >> m >> k;
	v.resize(5500);
	v2.resize(5500);
	v[1] = 1;
	int ans = 0;
	for (int i = 1; i<k; i++)
		iter();
	for (int i = 1; i<=5000; i++)
	{
		ans += eval(v[i],i);
		ans %= MOD;
	}
	cout << ans << '\n';
}