[UOJ#340][清华集训2017]小 Y 和恐怖的奴隶主（期望 DP + 矩阵乘法）

Address

洛谷P4007
UOJ#340
LOJ#2325

Solution

难道 $m$ 和 $k$ 特别小，容易想到一个状态：
$f[i][a][b][c]$ 表示当前血量为 $1$ 的随从有 $a$ 个，血量为 $2$ 的有 $b$ 个，血量为 $3$ 的有 $c$ 个的情况下，发动 $i$ 次攻击对 boss 造成的期望伤害。
边界：
$$f[0][a][b][c]=0$$
转移 1 ：攻击 boss 。
$$f[i][a][b][c]+=\frac{(1+f[i-1][a][b][c])}{1+a+b+c}$$
转移 2 ：击杀一个血量为 $1$ 的随从。
$$f[i][a][b][c]+=f[i-1][a-1][b][c]\times \frac{a}{1+a+b+c}$$
转移 3 ：随从数量达到上限时，攻击一个血量为 $2$ 的随从。
$$f[i][a][b][c]+=f[i-1][a+1][b-1][c]\times \frac{b}{1+a+b+c},a+b+c=k$$
转移 4 ：随从数量达到上限时，攻击一个血量为 $3$ 的随从。
$$f[i][a][b][c]+=f[i-1][a][b+1][c-1]\times \frac{c}{1+a+b+c},a+b+c=k$$
转移 5 ：随从数量未达到上限时，攻击一个血量为 $2$ 的随从，召唤一个血量为 $m$ 的随从。
（1） $m=2$ ：
$$f[i][a][b][c]+=f[i-1][a+1][b][c]\times \frac{b}{1+a+b+c}$$
（2） $m=3$ ：
$$f[i][a][b][c]+=f[i-1][a+1][b-1][c+1]\times \frac{b}{1+a+b+c}$$
转移 6 ：随从数量未达到上限时，攻击一个血量为 $3$ 的随从，召唤一个血量为 $m$ 的随从。
(1) m=3 ：
$$f[i][a][b][c]+=f[i-1][a][b+1][c]\times \frac{c}{1+a+b+c}$$
转移 5 和转移 6 的条件为 $a+b+c<k$ 。
注意到 $n$ 非常大，递推式是线性的，多次转移时系数不变，可以使用矩阵乘法优化 DP 。
注意到满足 $a+b+c\le k$ 的有序三元组 $(a,b,c)$ 最多有 $\left(\genfrac{}{}{0ex}{}{k+3}{3}\right)$ 个，
当 $k=8$ 时为 $165$ 。
可以矩乘。
最后一个小问题就是多组数据。
我们知道，矩阵快速幂优化 DP 的本质是倍增。
所以预处理出转移矩阵的 $2^0$ 、 $2^1$ 、 $2^2$ ， … 次幂。
然后每次询问就相当与一个初始向量乘上不超过 $O(\log n)$ 个转移矩阵。
复杂度 $$O({\left(\genfrac{}{}{0ex}{}{k+3}{3}\right)}^3\log n+T{\left(\genfrac{}{}{0ex}{}{k+3}{3}\right)}^2\log n)$$
一个卡常的技巧是取一个 $2\times 10^{18}$ 左右的， $998244353$ 的倍数 $R$ ，矩阵乘法时用 long long 加起来，当超过 $R$ 时减去 $R$ ，加完之后再取模。

Code

#include <cmath>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <algorithm>
#define For(i, a, b) for (i = a; i <= b; i++)

inline int read()
{
	int res = 0; bool bo = 0; char c;
	while (((c = getchar()) < '0' || c > '9') && c != '-');
	if (c == '-') bo = 1; else res = c - 48;
	while ((c = getchar()) >= '0' && c <= '9')
		res = (res << 3) + (res << 1) + (c - 48);
	return bo ? ~res + 1 : res;
}

typedef long long ll;

const int N = 170, M = 74, P = 10, ZZQ = 998244353;

int T, m, K, st[P][P][P], nn, inv[P];
ll n;

struct matr
{
	int n, m, a[N][N];
	
	matr() {}
	matr(int _n, int _m) : n(_n), m(_m)
	{
		memset(a, 0, sizeof(a));
	}
	
	friend inline matr operator * (const matr &a, const matr &b)
	{
		int i, j, k;
		matr res = matr(a.n, b.m);
		For (i, 1, res.n) For (j, 1, res.m)
		{
			ll sum = 0;
			For (k, 1, a.m)
			{
				sum += 1ll * a.a[i][k] * b.a[k][j];
				if (sum >= 2096313141300000000ll)
					sum -= 2096313141300000000ll;
			}
			res.a[i][j] = sum % ZZQ;
		}
		return res;
	}
} pw[M], St, Ans;

int main()
{
	int i, j, k;
	T = read(); m = read(); K = read();
	inv[1] = 1;
	For (i, 2, K + 1) inv[i] = 1ll * (ZZQ - ZZQ / i) * inv[ZZQ % i] % ZZQ;
	For (i, 0, K) For (j, 0, K) For (k, 0, K)
		if (i + j + k <= K) st[i][j][k] = ++nn;
	pw[0] = matr(nn + 1, nn + 1);
	St = matr(nn + 1, 1);
	St.a[nn + 1][1] = 1;
	For (i, 0, K) For (j, 0, K) For (k, 0, K)
	{
		if (i + j + k > K) continue;
		int x = st[i][j][k], tmp = inv[i + j + k + 1];
		if (i) pw[0].a[x][st[i - 1][j][k]] = 1ll * i * tmp % ZZQ;
		if (j)
		{
			if (i + j + k == K)
				pw[0].a[x][st[i + 1][j - 1][k]] = 1ll * j * tmp % ZZQ;
			else if (m == 1)
				pw[0].a[x][st[i + 2][j - 1][k]] = 1ll * j * tmp % ZZQ;
			else if (m == 2)
				pw[0].a[x][st[i + 1][j][k]] = 1ll * j * tmp % ZZQ;
			else pw[0].a[x][st[i + 1][j - 1][k + 1]] = 1ll * j * tmp % ZZQ;
		}
		if (k)
		{
			if (i + j + k == K)
				pw[0].a[x][st[i][j + 1][k - 1]] = 1ll * k * tmp % ZZQ;
			else if (m == 1)
				pw[0].a[x][st[i + 1][j + 1][k - 1]] = 1ll * k * tmp % ZZQ;
			else if (m == 2)
				pw[0].a[x][st[i][j + 2][k - 1]] = 1ll * k * tmp % ZZQ;
			else pw[0].a[x][st[i][j + 1][k]] = 1ll * k * tmp % ZZQ;
		}
		pw[0].a[x][x] = pw[0].a[x][nn + 1] = tmp;
	}
	pw[0].a[nn + 1][nn + 1] = 1;
	For (i, 1, 59) pw[i] = pw[i - 1] * pw[i - 1];
	while (T--)
	{
		scanf("%lld", &n);
		Ans = St;
		For (i, 0, 59)
			if ((n >> i) & 1) Ans = pw[i] * Ans;
		if (m == 1) printf("%d\n", Ans.a[st[1][0][0]][1]);
		else if (m == 2) printf("%d\n", Ans.a[st[0][1][0]][1]);
		else printf("%d\n", Ans.a[st[0][0][1]][1]);
	}
	return 0;
}