LOJ 2143「SHOI2017」组合数问题

题目链接：LOJ 2143

好像是循环矩阵一样的东西。
今天早上才做过一道用矩阵的[题目][https://www.luogu.com.cn/problem/P2886]

按照题目去算显然是会超时的， 可以按\(\bmod k\)的余数分类。
设\(f[i][j]\)表示从\(i\)个数中选若干个数，余数为\(j\)的方案数。
转移是\(f[i][j]=f[i][j]+f[x][j-1]*f[y][j-2]\),\(((j-1+j-2)\bmod k=j)\).
因为只和\(j\)有关， 所以可以用矩阵快速幂优化。
注意\(k\)可能等于1，初始化的时候注意一下。
时间复杂度\(\mathjaxcal{O}(k^{2\log {nk}})\)。

#include <bits/stdc++.h>
using namespace std;

int n, k, r, p;

struct Martix {
	int a[55];
	Martix operator * (const Martix& x) const {
		Martix tmp;
		memset(tmp.a, 0, sizeof tmp.a);
		for (int i = 0; i < k; i++)
			for (int j = 0; j < k; j++) {
                tmp.a[(i + j) % k] = (tmp.a[(i + j) % k] + a[i] * 1ll * x.a[j] % p) % p;
			}
		return tmp;
	}
} ans;
Martix fpow(Martix a, long long b) {
	Martix ans;
	memset(ans.a, 0, sizeof ans.a);
    ans.a[0] = 1;
    for (; b; b >>= 1, a = a * a) if (b & 1) ans = a * ans;
    return ans;
}

int main() {
    scanf("%d %d %d %d", &n, &p, &k, &r);
    ans.a[1 % k]++, ans.a[0]++;
    ans = fpow(ans, n * 1ll * k);
    printf("%d\n", ans.a[r]);
    return 0;
}