洛谷P4007 小 Y 和恐怖的奴隶主(期望dp 矩阵乘法)

题意

题目链接

Sol

首先不难想到一种暴力dp，设\(f[i][a][b][c]\)表示还有\(i\)轮没打，场上有\(a\)个1血，\(b\)个2血，\(c\)个三血

发现状态数只有\(s = 166\)个，复杂度为\(O(ns)\)

矩乘优化一下复杂度为\(O(s^3 logn T)\)，还是过不去。

因为每次询问都是独立的，那么可以预处理出\(2^i\)的转移矩阵，回答询问只需要拿一个行向量去乘log个矩阵

构造矩阵的时候可以加一个列向量表示期望

#include<bits/stdc++.h>
#define LL long long
using namespace std;
const int B = 60, mod = 998244353;
template <typename A, typename B> inline bool chmin(A &a, B b){if(a > b) {a = b; return 1;} return 0;}
template <typename A, typename B> inline bool chmax(A &a, B b){if(a < b) {a = b; return 1;} return 0;}
template <typename A, typename B> inline LL add(A x, B y) {if(x + y < 0) return x + y + mod; return x + y >= mod ? x + y - mod : x + y;}
template <typename A, typename B> inline void add2(A &x, B y) {if(x + y < 0) x = x + y + mod; else x = (x + y >= mod ? x + y - mod : x + y);}
LL mul(int x, int y) {return 1ll * x * y % mod;}

inline LL read() {
	char c = getchar(); LL x = 0, f = 1;
	while(c < '0' || c > '9') {if(c == '-') f = -1; c = getchar();}
	while(c >= '0' && c <= '9') x = x * 10 + c - '0', c = getchar();
	return x * f;
}
int fp(int a, int p) {
	int base = 1;
	while(p) {
		if(p & 1) base = mul(base, a);
		a = mul(a, a); p >>= 1;
	}
	return base;
}
int T, M, K;

namespace S3 {
	int id[11][11][11], cnt, Lim;
	int ans[168];
	LL inv[11];

	struct Ma {
		int m[168][168];
		Ma() {
			memset(m, 0, sizeof(m));
		}
		void init() {
			for(int i = 0; i <= Lim; i++) m[i][i] = 1;
		}
		void print() {
			for(int i = 1; i <= Lim; i++, puts(""))
				for(int j = 1; j <= Lim; j++)
					printf("%d ", m[i][j]);
		}
		Ma operator * (const Ma &rhs) const {
			Ma gg = {};
			for(int i = 1; i <= Lim; i++)
				for(int j = 1; j <= Lim; j++) {
					__int128 tmp = 0;
					for(int k = 1; k <= Lim; k++)
						tmp += mul(m[i][k], rhs.m[k][j]);
					tmp %= mod;
					gg.m[i][j] = tmp;
				}

			return gg;
		}
	}f[B + 1];
	void Pre() {
		for(int i = 1; i <= K + 1; i++) inv[i] = fp(i, mod - 2);
		for(int a = 0; a <= K; a++)
			for(int b = 0; a + b <= K; b++)
				for(int c = 0; a + b + c <= K; c++)
					id[a][b][c] = ++cnt;
		for(int a = 0; a <= K; a++)
			for(int b = 0; a + b <= K; b++)
				for(int c = 0; a + b + c <= K; c++) {
					int down = inv[a + b + c + 1], tag = (a + b + c < K), now = id[a][b][c];
					if(a) f[0].m[now][id[a - 1][b][c]] = mul(a, down);
					if(b) f[0].m[now][id[a + 1][b - 1][c + tag]] = mul(b, down);
					if(c) f[0].m[now][id[a][b + 1][c - 1 + tag]] = mul(c, down);
					f[0].m[now][now] = down;
					f[0].m[now][cnt + 1] = down;
				}
		f[0].m[cnt + 1][cnt + 1] = 1;
		Lim = cnt + 1;
		for(int i = 1; i <= B; i++) f[i] = f[i - 1] * f[i - 1];
	}
	int tmp[168];
	void mul(Ma a) {
		memset(tmp, 0, sizeof(tmp));
		for(int j = 1; j <= Lim; j++)
			for(int i = 1; i <= Lim; i++)
				add2(tmp[j], 1ll * ans[i] * a.m[i][j] % mod);
		memcpy(ans, tmp, sizeof(tmp));
	}
	void MatrixPow(LL p) {
		for(int i = 0; p; p >>= 1, i++)
			if(p & 1)
				mul(f[i]);
	}
	void work() {
		Pre();
		while(T--) {
			LL n = read();
			memset(ans, 0, sizeof(ans)); ans[id[0][0][1]] = 1;
			MatrixPow(n);
			cout << ans[cnt + 1] << '\n';
		}
	}
}

namespace S2 {
	int id[11][11], cnt, Lim;
	int ans[168];
	LL inv[11];

	struct Ma {
		int m[168][168];
		Ma() {
			memset(m, 0, sizeof(m));
		}
		void init() {
			for(int i = 0; i <= Lim; i++) m[i][i] = 1;
		}
		void print() {
			for(int i = 1; i <= Lim; i++, puts(""))
				for(int j = 1; j <= Lim; j++)
					printf("%d ", m[i][j]);
		}
		Ma operator * (const Ma &rhs) const {
			Ma gg = {};
			for(int i = 1; i <= Lim; i++)
				for(int j = 1; j <= Lim; j++) {
					__int128 tmp = 0;
					for(int k = 1; k <= Lim; k++)
						tmp += mul(m[i][k], rhs.m[k][j]);
					tmp %= mod;
					gg.m[i][j] = tmp;
				}

			return gg;
		}
	}f[B + 1];
	void Pre() {
		for(int i = 1; i <= K + 1; i++) inv[i] = fp(i, mod - 2);
		for(int a = 0; a <= K; a++)
			for(int b = 0; a + b <= K; b++)
				id[a][b] = ++cnt;
		for(int a = 0; a <= K; a++)
			for(int b = 0; a + b <= K; b++) {
				int down = inv[a + b + 1], tag = (a + b < K), now = id[a][b];
				if(a) f[0].m[now][id[a - 1][b]] = mul(a, down);
				if(b) f[0].m[now][id[a + 1][b - 1 + tag]] = mul(b, down);
				f[0].m[now][now] = down;
				f[0].m[now][cnt + 1] = down;
			}
		f[0].m[cnt + 1][cnt + 1] = 1;
		Lim = cnt + 1;
		for(int i = 1; i <= B; i++) f[i] = f[i - 1] * f[i - 1];
	}
	int tmp[168];
	void mul(Ma a) {
		memset(tmp, 0, sizeof(tmp));
		for(int j = 1; j <= Lim; j++)
			for(int i = 1; i <= Lim; i++)
				add2(tmp[j], 1ll * ans[i] * a.m[i][j] % mod);
		memcpy(ans, tmp, sizeof(tmp));
	}
	void MatrixPow(LL p) {
		for(int i = 0; p; p >>= 1, i++)
			if(p & 1)
				mul(f[i]);
	}
	void work() {
		Pre();
		while(T--) {
			LL n = read();
			memset(ans, 0, sizeof(ans)); ans[id[0][1]] = 1;
			MatrixPow(n);
			cout << ans[cnt + 1] << '\n';
		}
	}
}

namespace S1 {
	int N,  f[12][9][9][9];
	int inv(int a) {
		return fp(a, mod - 2);
	}
	void work() {
		N = 11;
		for(int i = 1; i <= N; i++) {
			for(int a = 0; a <= K; a++) {
				for(int b = 0; a + b <= K; b++) {
					for(int c = 0; a + b + c <= K; c++) {
						int down = a + b + c + 1;
						if(a) add2(f[i][a][b][c], mul(mul(a, inv(down)), f[i - 1][a - 1][b][c]));
						if(b) {
							if(down <= K) add2(f[i][a][b][c], mul(mul(b, inv(down)), f[i - 1][a + 1][b - 1 + (M == 2)][c + (M == 3)]));
							else add2(f[i][a][b][c], mul(mul(b, inv(down)), f[i - 1][a + 1][b - 1][c]));
						}
						if(c) {
							if(down <= K) add2(f[i][a][b][c], mul(mul(c, inv(down)), f[i - 1][a][b + 1 + (M == 2)][c - 1 + (M == 3)]));
							else add2(f[i][a][b][c], mul(mul(c, inv(down)), f[i - 1][a][b + 1][c - 1]));
						}
						add2(f[i][a][b][c], mul(inv(down), f[i - 1][a][b][c] + 1));
					}
				}
			}
		}
		while(T--) {
			int n = read();
			printf("%d\n", f[n][M == 1][M == 2][M == 3]);
		}
	}
}

int main() {
	T = read(); M = read(); K = read();
	if(M == 1) S1::work();
	else if(M == 2) S2::work();
	else S3::work();

	return 0;
}