[清华集训 2017]小 Y 和恐怖的奴隶主 [概率DP+矩阵乘法]

传送门

我开始想把8个随从的血压成一个状态...

其实不然, 我们不需要知道每个随从的血的情况, 我们只需要知道 血为 1/ 2/ 3的个数

 表示有 i 个血为1, j 个血为2, k 个血为 3, 还剩 now 次攻击次数打到的期望

我们把3维压成一维, 大概 165 的状态, 我们预处理 j 状态 可以转移到 i 状态, 然后矩阵乘法

考虑后面那个加一怎么办, 我们加一维表示 1 / (i+j+k+1), 然后向量的最后一维设成1, 每次就可以加上贡献了

但是 165 ^ 3 * 500 要 T, 我们发现向量 * 矩阵是 O(n^2) 的, 于是预处理2的整次幂的矩阵

复杂度就是 O(165 ^ 3 * log(n) + 165 ^ 2 * 500)

---

#include<bits/stdc++.h>
#define N 200
#define R register
using namespace std;
typedef long long ll;
const int Mod = 998244353;
int T; ll n; int m, K;
int pos[10][10][10], x[N], y[N], z[N], tot, inv[N];
ll add(ll a, ll b){ return (a+b) % Mod;}
ll mul(ll a, ll b){ return (a*b) % Mod;}
ll power(ll a, ll b){ ll ans = 1; for(;b;b>>=1){if(b&1) ans = mul(ans, a); a = mul(a, a);} return ans;}
struct Matrix{
	ll a[N][N];
	Matrix(){ memset(a, 0, sizeof(a));}
	inline friend Matrix operator * (const Matrix &A, const Matrix &B){
		Matrix C; for(R int i=1; i<=tot; i++)
		for(R int j=1; j<=tot; j++) for(R int k=1; k<=tot; k++)
			C.a[i][j] = add(C.a[i][j], mul(A.a[i][k], B.a[k][j]));
		return C;
	} 
	inline friend Matrix operator ^ (const Matrix &A, const Matrix &B){
		Matrix C; for(R int j=1; j<=tot; j++)
		for(R int k=1; k<=tot; k++) C.a[1][j] = add(C.a[1][j], mul(A.a[1][k], B.a[k][j]));
		return C;
	}
}F[70];
int main(){
	scanf("%d%d%d", &T, &m, &K);
	for(int i=1; i<=20; i++) inv[i] = power(i, Mod - 2);
	for(int i=0; i<=K; i++){
		for(int j=0; j<=K-i; j++){
			for(int k=0; k<=K-i-j; k++){
				if(m == 1 && (j || k)) continue;
				if(m == 2 && k) continue;
				pos[i][j][k] = ++tot;
				x[tot] = i; y[tot] = j; z[tot] = k;
			}
		}
	} tot++;
	for(int i=1; i<tot; i++) F[0].a[tot][i] = inv[x[i] + y[i] + z[i] + 1];
	F[0].a[tot][tot] = 1;
	for(int i=1; i<tot; i++){
		for(int j=1; j<tot; j++){
			if(i == j) F[0].a[i][i] = inv[x[i] + y[i] + z[i] + 1];
			else{
				if(z[i]){
					int xx = x[i], yy = y[i] + 1, zz = z[i] - 1;
					if(xx + yy + zz < K){ xx += (m == 1); yy += (m == 2); zz += (m == 3);}
					if(xx == x[j] && yy == y[j] && zz == z[j]) 
						F[0].a[j][i] = mul(z[i], inv[x[i] + y[i] + z[i] + 1]); 
				}
				if(y[i]){
					int xx = x[i] + 1, yy = y[i] - 1, zz = z[i];
					if(xx + yy + zz < K){ xx += (m == 1); yy += (m == 2); zz += (m == 3);}
					if(xx == x[j] && yy == y[j] && zz == z[j]) 
						F[0].a[j][i] = mul(y[i], inv[x[i] + y[i] + z[i] + 1]);
				}
				if(x[i]){
					int xx = x[i] - 1, yy = y[i], zz = z[i];
					if(xx == x[j] && yy == y[j] && zz == z[j])
						F[0].a[j][i] = mul(x[i], inv[x[i] + y[i] + z[i] + 1]);
				}
			}
		}
	}
	for(int i=1; i<=60; i++) F[i] = F[i-1] * F[i-1];
	while(T--){
		scanf("%lld", &n); 
		Matrix ans; ans.a[1][tot] = 1;
		for(int i=60; i>=0; i--){
			if(n >= (1ll << i)){n -= (1ll<<i); ans = ans ^ F[i];}
		} printf("%lld\n", ans.a[1][pos[m == 1][m == 2][m == 3]]);
	} return 0;
}