LOJ2143. 「SHOI2017」组合数问题

#2143. 「SHOI2017」组合数问题

求 $\left( \sum_{i = 0}^\infty \mathrm{C}_{nk}^{ik + r} \right) \bmod p$ 。

处理一下式子，写成 $\sum _ {i = 1} ^ {nk} \binom {nk} {i} [i \equiv r \pmod p]$ ，即从 $nk$ 里面选择 $i$ 个物品，其中 $i$ 满足一定条件。那么不妨依此考虑每一个物品是否被选择，记 $f _ {i, j}$ 表示当前考虑第 $i$ 个物品，选择了 $j$ 个数，而 $j$ 只需考虑模 $k$ 的余数，有转移 $f _ {i, j} = f _ {i - 1, j - 1} + f _ {i - 1, j}$ 。而 $nk$ 很大，可以矩阵优化，复杂度为 $O(n ^ 3 \log k)$ 。

查看代码

#include <cstdio>
#include <iostream>
using namespace std;
typedef long long L;
const int N = 60;
int p;
struct Matrix
{
    int n, m, w[N][N];
    Matrix (int x, int y, int k)
    {
        n = x, m = y;
        for (int i = 0; i < n; ++i)
            for (int j = 0; j < m; ++j)
                w[i][j] = i == j ? k : 0;
    }
    friend Matrix operator* (Matrix x, Matrix y)
    {
        Matrix res(x.n, y.m, 0);
        for (int i = 0; i < res.n; ++i)
            for (int j = 0; j < res.m; ++j)
                for (int k = 0; k < x.m; ++k)
                    res.w[i][j] = (res.w[i][j] + (L)x.w[i][k] * y.w[k][j]) % p;
        return res;
    }
    friend Matrix operator^ (Matrix b, L k)
    {
        Matrix res(b.n, b.m, 1);
        for (; k; k >>= 1, b = b * b)
            if (k & 1) res = res * b;
        return res;
    }
};
int main ()
{
    int n, k, r;
    cin >> n >> p >> k >> r;
    Matrix A(1, k, 0), B(k, k, 0);
    A.w[0][0] = 1;
    for (int i = 0; i < k; ++i)
        ++B.w[i][i], ++B.w[i][(i + 1) % k];
    cout << (A * (B ^ (L)n * k)).w[0][r];
    return 0;
}