BOJ 33688 불의 군주 라그나로스 2

문제 링크입니다.

문제 요약

라그나로스 $N$체가 내 필드에 있습니다. 적 영웅의 체력과 하수인이 주어졌을 때, 턴 종료 시 승리하는 경우의 수를 구하세요.

풀이

라그나로스의 피해량, 적 영웅과 하수인의 체력을 버리고 대신 각 개체를 처치하기 위해 필요한 공격 횟수를 들고 옵니다. 라그나로스를 공으로, 각 개체를 상자로 치환해 조합론적 모델로 해석합니다.

상자는 $0$번부터 $M$번까지 총 $M+1$개 있습니다. 각 상자는 용량 $C_i$가 정해져 있어, 용량을 초과하는 공은 집어넣을 수 없습니다. 공은 $1$번부터 $N$번까지 순서대로 자리가 남아 있는 상자 중 하나에 집어넣습니다. $0$번 상자에 공이 가득 차는 순간 행동을 종료하며, 이때 $0$번 상자에 공이 가득 차는 경우의 수를 구하면 됩니다.

가능한 모든 $n$에 대해 $n$번째 공으로 $0$번 상자에 공이 가득 차는 경우의 수를 더하면 되겠네요. 이를 위해 $0$번 상자의 용량을 1 깎고, 공의 순서를 따져야 하니 지수생성함수를 고려합니다.

• $0$번 상자의 경우 반드시 $C_0 - 1$개의 공이 들어차야 하므로, 생성함수는 $\frac{1}{(C_0 - 1)!}x^{C_0-1}$입니다.
• $i$번 상자($1 \le i \le M$)의 경우 그냥 $\sum\limits_{t=0}^{C_i}{\frac{1}{t!}x^t}$가 생성함수입니다.

이제 이 다항식들을 싹 다 곱하면, $t$차항의 계수 $D_t$에 대해 $t! D_t$가 $t+1$번째 공으로 $0$번 상자를 채우고 끝내는 경우의 수가 됩니다.

다항식을 곱하는 건 차수가 작은 것끼리 곱하면서 필요 이상 차수의 항은 잘라내는 식으로 하면 좀 덜 잘 짜도 간당간당하게 시간 제한 안에 들어갑니다.

코드

#include<bits/stdc++.h>
using namespace std;
using ll = long long;
using pii = pair<int, int>;
using pll = pair<ll, ll>;
const ll mod = 998244353;
const ll w = 3;
int NMAX = 50010;

ll pw(ll n, ll r){
	ll ret = 1;
	while(r){
		if(r&1) ret = ret*n%mod;
		n = n*n%mod;
		r >>= 1;
	}
	return ret;
}

inline ll inv(ll n){return pw(n, mod-2);}

vector<ll> p[50505];
ll fac[50505], ifac[50505];

void ntt(vector<ll>& f, bool flag = 0){
    int n = f.size(), j = 0;
    vector<ll> root(n>>1);
    for(int i = 1; i < n; i++){
        int bit = n>>1;
        while(j >= bit){
            j -= bit; bit >>= 1;
        }
        j += bit;
        if(i < j) swap(f[i], f[j]);
    }
    ll ang = pw(w, (mod - 1) / n); if(flag) ang = inv(ang);
    root[0] = 1; for(int i=1; i<(n >> 1); i++) root[i] = root[i-1] * ang % mod;
    for(int i=2; i<=n; i<<=1){
        int step = n / i;
        for(int j=0; j<n; j+=i){
            for(int k=0; k<(i >> 1); k++){
                ll u = f[j | k], v = f[j | k | i >> 1] * root[step * k] % mod;
                f[j | k] = (u + v) % mod;
                f[j | k | i >> 1] = (u - v) % mod;
                if(f[j | k | i >> 1] < 0) f[j | k | i >> 1] += mod;
            }
        }
    }
    ll t = inv(n);
    if(flag) for(int i=0; i<n; i++) f[i] = f[i] * t % mod;
}

void mult(vector<ll>& a, vector<ll>& b){
	int n = 2, N = a.size()+b.size()-1; while(n < a.size()+b.size()) n <<= 1;
	a.resize(n); b.resize(n); ntt(a); ntt(b);
	for(int i = 0; i < n; i++) a[i] = a[i]*b[i]%mod;
	ntt(a, 1); a.resize(min(N, NMAX)); return;
}

priority_queue<pii> pq;

int main(){
	cin.tie(0); cout.tie(0); ios::sync_with_stdio(false);
	fac[0] = 1; for(int i = 1; i <= 50500; i++) fac[i] = fac[i-1]*i%mod;
	ifac[50500] = pw(fac[50500], mod-2); for(int i = 50500; i; i--) ifac[i-1] = ifac[i]*i%mod;
	int n, m, x, y; cin >> n >> m >> x >> y; y = (y-1)/x+1; p[0].resize(y); p[0][y-1] = ifac[y-1]; pq.emplace(-y, 0);
	for(int i = 1; i <= m; i++){
		int h; cin >> h; h = (h-1)/x+1; for(int j = 0; j <= h; j++) p[i].push_back(ifac[j]); pq.emplace(-h-1, i);
	}
	while(pq.size() > 1){
		auto [d, i] = pq.top(); pq.pop();
		auto [dd, j] = pq.top(); pq.pop();
		mult(p[i], p[j]); pq.emplace(-p[i].size(), i);
	}
	ll ans = 0; int idx = pq.top().second; n = min(n, (int)p[idx].size());
	for(int i = y-1; i < n; i++) ans = (ans+fac[i]*p[idx][i]%mod)%mod;
	cout << ans;
}