[题解]AT_abc256_g [ABC256G] Black and White Stones

思路

容易看出来是个 DP 题，但是你发现 DP 的起点是不好确定的，于是假定第一条边的起点是黑色。然后你发现设为白色的贡献与黑色是相同的，于是直接令第一条边的起点是黑色，最后答案乘以 $2$ 即可。

然后就可以愉快的 DP 了。

首先枚举每条边白色点的数量 $k$，定义 $dp_{i,0/1}$ 表示确定了前 $i$ 条边，且第 $i$ 条边的终点为 黑色/白色。转移就比较显然了，递推起点为 $dp_{0,0} = 1$：

$$\left\{\begin{matrix} dp_{i,0} = dp_{i - 1,0} \times \binom{d - 1}{k} + dp_{i - 1,1} \times \binom{d - 1}{k - 1}\\ dp_{i,1} = dp_{i - 1,0} \times \binom{d - 1}{k - 1} + dp_{i - 1,1} \times \binom{d - 1}{k - 2} \end{matrix}\right.$$

但是这样转移的复杂度高达 $\Theta(nd)$，但是仔细观察发现这个式子与矩阵乘法类似。

令 $c_1 = \binom{d - 1}{k},c_2 = \binom{d - 1}{k - 1},c_3 = \binom{d - 1}{k - 2}$。

于是考虑维护 $\begin{bmatrix} dp_{i,0} & dp_{i,1} \end{bmatrix}$ 这个矩阵，它将转移为 $\begin{bmatrix} dp_{i,0} \times c_1 + dp_{i,1} \times c_2 & dp_{i,0} \times c_2 + dp_{i,1} \times c_3 \end{bmatrix}$。

容易构造出转移矩阵：

$$\begin{bmatrix} c_1 & c_2\\ c_2 & c_3 \end{bmatrix}$$

直接矩阵快速幂优化即可。

Code

#include <bits/stdc++.h>
#define re register
#define int long long
#define Add(a,b) (((a) % mod + (b) % mod) % mod)
#define Mul(a,b) (((a) % mod) * ((b) % mod) % mod)

using namespace std;

const int N = 1e7 + 10,M = 1e4 + 10,mod = 998244353;
int n,d,ans;
int fac[M],inv[M];

struct mat{
	int n,m;
	int mt[5][5];
	
	mat(int a,int b){
		n = a;
		m = b;
		memset(mt,0,sizeof(mt));
	}

    mat friend operator *(const mat &a,const mat &b){
        mat t(a.m,b.m);
        for (re int i = 1;i <= a.n;i++){
            for (re int j = 1;j <= b.m;j++){
                for (re int k = 1;k <= a.m;k++) t.mt[i][j] = Add(t.mt[i][j],Mul(a.mt[i][k],b.mt[k][j]));
            }
        }
        return t;
    }
}base(2,2),dp(1,2);

inline int read(){
    int r = 0,w = 1;
    char c = getchar();
    while (c < '0' || c > '9'){
        if (c == '-') w = -1;
        c = getchar();
    }
    while (c >= '0' && c <= '9'){
        r = (r << 3) + (r << 1) + (c ^ 48);
        c = getchar();
    }
    return r * w;
}

inline int qmival(int a,int b){
    int res = 1;
    while (b){
        if (b & 1) res = Mul(res,a);
        a = Mul(a,a); b >>= 1;
    }
    return res;
}

inline mat qmi(mat a,int b){
	mat c(a.n,a.n);
	for (re int i = 1;i <= c.n;i++) c.mt[i][i] = 1;
	while (b){
		if (b & 1) c = c * a;
		a = a * a;
		b >>= 1;
	}
	return c;
}

inline void init(){
    fac[0] = 1;
    for (re int i = 1;i <= d;i++) fac[i] = Mul(fac[i - 1],i);
    inv[d] = qmival(fac[d],mod - 2);
    for (re int i = d - 1;~i;i--) inv[i] = Mul(inv[i + 1],i + 1);
}

inline int C(int n,int m){
    if (n < m || m < 0) return 0;
    return Mul(fac[n],Mul(inv[n - m],inv[m]));
}

inline int f(int k){
    int c1 = C(d - 1,k),c2 = C(d - 1,k - 1),c3 = C(d - 1,k - 2);
    dp.mt[1][1] = 1;
    base.mt[1][1] = c1; base.mt[1][2] = c2; base.mt[2][1] = c2; base.mt[2][2] = c3;
    mat ans = dp * qmi(base,n);
    return Mul(ans.mt[1][1],2);
}

signed main(){
    n = read(),d = read(); init();
    for (re int i = 0;i <= d;i++) ans = Add(ans,f(i));
    printf("%lld",ans);
    return 0;
}