UOJ340 【清华集训2017】小 Y 和恐怖的奴隶主题解

题目链接：#340. 【清华集训2017】小 Y 和恐怖的奴隶主

题意：

小 Y 是一个喜欢玩游戏的 OIer。一天，她正在玩一款游戏，要打一个 Boss。

虽然这个 Boss 有 \(10^{100}\) 点生命值，但它只带了一个随从——一个只有 \(m\) 点生命值的 “恐怖的奴隶主”。

这个 “恐怖的奴隶主” 有一个特殊的技能：每当它被扣减生命值但没有死亡（死亡即生命值 \(\leq 0\)），且 Boss 的随从数量小于上限 \(k\)，便会召唤一个新的具有 \(m\) 点生命值的 “恐怖的奴隶主”。

现在小 Y 可以进行 \(n\) 次攻击，每次攻击时，会从 Boss 以及 Boss 的所有随从中的等概率随机选择一个，并扣减 \(1\) 点生命值，她想知道进行 \(n\) 次攻击后扣减 Boss 的生命值点数的期望。为了避免精度误差，你的答案需要对 \(998244353\) 取模。

输入格式：

从标准输入读入数据。

输入第一行包含三个正整数 \(T, m, k\)，\(T\) 表示询问组数，\(m, k\) 的含义见题目描述。

接下来 \(T\) 行，每行包含一个正整数 \(n\)，表示询问进行 \(n\) 次攻击后扣减 Boss 的生命值点数的期望。

输出格式：

输出共 \(T\) 行，对于每个询问输出一行一个非负整数，表示该询问的答案对 \(998244353\) 取模的结果。

可以证明，所求期望一定是一个有理数，设其为 \(\frac{p}{q}\)（\(\mathrm{gcd}(p,q) = 1\)），那么你输出的数 \(x\) 要满足 \(p \equiv qx \pmod{998244353}\)。

数据范围：

\(1 \le T \le 1000,~1 \le n \le 10^{18},~m = 1,2,3,~1\le k \le 8\)

~~据说概率要顺推，期望要逆推？~~

容易想到是矩阵快速幂优化期望dp。

注意到 \(m\) 很小，且 \(k\) 也很小，所以我们直接枚举每个血量各有多少人。

设 \(f_{i,a,b,c}\) 表示前 \(i\) 次攻击后，场上有 \(a\) 个 \(1\) 点血，\(b\) 个 \(2\) 点血，\(c\) 个 \(3\) 点血的期望。

转移很简单，就是乘一乘概率系数啥的，这里就略去了。~~或者可以看参考文献[2]的分析~~

注意到 \(f_i\) 只有 \(165 + 1\) 个状态，其中一个是额外加上的，表示 boss 扣血的期望 \(g\) 。

关于构建矩阵，估计也有很多人和我一样知道矩阵快速幂，但不会代码实现矩阵的构建。

最开始学斐波那契数列的矩阵快速幂方法时，只是看这个矩阵挺好用，怎么正好就推出来下一个。

现在做了这道题，感觉对这个矩阵的理解又加深了一些。

对于转移矩阵 \(M\) ，考察 \(M_{i,j}\) 的数学意义。

状态数组（行向量）左乘 \(M_{i,j}\) 其实就是：「原向量第 \(i\) 个状态」转移到「新向量第 \(j\) 个状态」时乘的系数。

也可以说，「新向量第 \(j\) 个状态」由「原向量第 \(i\) 个状态」转移过来时乘的系数。

举个例子，斐波那契数列 \[ \begin{bmatrix} f_i&f_{i-1} \end{bmatrix} \times \begin{bmatrix} 1&1\\1&0 \end{bmatrix} = \begin{bmatrix} f_{i+1}&f_{i} \end{bmatrix} \] 因为有递推式 \(f_1=1,~f_2=1,~f_i = f_{i-1} + f_{i-2}\) 。

本质上矩阵就是多个行向量构成的，因此对于矩阵乘法，我们也可以用这样的方式理解，如 \[ \begin{bmatrix}1&2\\3&4\end{bmatrix} \times \begin{bmatrix}5&6\\7&8\end{bmatrix} = \begin{bmatrix}5 \times 1 + 7 \times 2 & 6 \times 1+ 8 \times 2 \\ 5 \times 3 + 7 \times 4&6 \times 3 + 8 \times 4\end{bmatrix} = \begin{bmatrix}19&22\\43&50\end{bmatrix} \] 属于是复习一下矩阵快速幂了（

然而直接矩阵快速幂是 \(\mathcal{O}(T166^3\log n)\) 的，会 TLE。

考虑预处理 \(2^k\) 次幂的矩阵，这样查询就可以少掉一个 \(M\times M\) 的第 \(k\) 维枚举了。

时间复杂度 \(\mathcal{O}(T166^2\log n)\)

代码：（这题卡常，可以用 __int128 优化取模操作）

#include <bits/stdc++.h>
using namespace std;
// #define int long long
// #define INF 0x3f3f3f3f3f3f3f3f
#pragma GCC optimize("Ofast,unroll-loops,fast-math")
typedef long long ll;
const ll mod = 998244353;
void up(int &x,int y) { x < y ? x = y : 0; }
void down(int &x,int y) { x > y ? x = y : 0; }
typedef vector< vector<int> > mat;
#define N ((int)(166 + 15))
#define K 11

mat M[70],ans;
int n,Q,m,k,cnt,inv[10],id[K][K][K];
int qpow(int a,int b)
{
    int ans = 1, base = a % mod;
    for(; b; b >>= 1)
    {
        if(b & 1) ans = 1ll * ans * base % mod;
        base = 1ll * base * base % mod;
    }
    return ans;
}
mat operator * (mat &A, mat &B)
{
    int n = A.size() - 1, m = A[0].size() - 1, p = B[0].size() - 1;
    mat C(n + 1, vector<int>(p + 1));
    for(int i=1; i<=n; i++)
        for(int j=1; j<=p; j++)
        {
            __int128 t = 0;
            for(int k=1; k<=m; k++) t += 1ll * A[i][k] * B[k][j];
            C[i][j] = t % mod;
    }
    return C;
}
signed main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);cout.tie(0);
    // freopen("check.in","r",stdin);
    // freopen("check.out","w",stdout);
    cin >> Q >> m >> k; M[0].resize(N, vector<int>(N));
    for(int i=1; i<=9; i++) inv[i] = qpow(i, mod - 2);
    for(int a=0; a<=k; a++)
        for(int b=0; b <= (m >= 2 ? k - a : 0); b++)
            for(int c=0; c <= (m == 3 ? k - a - b : 0); c++) id[a][b][c] = ++cnt;
    ++cnt; M[0][cnt][cnt] = 1; ans.resize(2, vector<int>(cnt + 1));
    for(int a=0; a<=k; a++)
        for(int b=0; b <= (m >= 2 ? k - a : 0); b++)
            for(int c=0; c <= (m == 3 ? k - a - b : 0); c++)
            {
                int i = id[a][b][c], iv = inv[a + b + c + 1], d = a + b + c < k;
                switch(m)
                {
                    case 1:
                        if(a) M[0][i][id[a - 1][b][c]] = 1ll * a * iv % mod;
                        break;
                    case 2:
                        if(a) M[0][i][id[a - 1][b][c]] = 1ll * a * iv % mod;
                        if(b) M[0][i][id[a + 1][b - 1 + d][c]] = 1ll * b * iv % mod;
                        break;
                    case 3:
                        if(a) M[0][i][id[a - 1][b][c]] = 1ll * a * iv % mod;
                        if(b) M[0][i][id[a + 1][b - 1][c + d]] = 1ll * b * iv % mod;
                        if(c) M[0][i][id[a][b + 1][c - 1 + d]] = 1ll * c * iv % mod;
                        break;
                }
                M[0][i][i] = M[0][i][cnt] = iv;
            }
    for(int i=1; i<=60; i++) M[i] = M[i-1] * M[i-1];
    for(ll n; Q; Q--)
    {
        cin >> n;
        for(int i=1; i<=cnt; i++) ans[1][i] = 0;
        ans[1][id[m == 1][m == 2][m == 3]] = 1;
        for(int i=0; i<=60; i++) if((n >> i) & 1) ans = ans * M[i];
        cout << ans[1][cnt] << '\n';
    }
    return 0;
}