【BZOJ4870】组合数问题（六省联考2017）-矩阵优化DP

测试地址：组合数问题
做法：这题在THUSC2017试机时看到了，当时觉得非常有趣，于是就找到了原题来做，感觉实在是妙啊……
这一题应该使用矩阵优化DP来解决。
第一眼看上去这题感觉非常难做，推式子又感觉推不出，但其实我们观察一下这个式子，发现它的意义可以表达成：从nk$nk$个物品里取出的物品数量对k$k$取模结果为r$r$的方案数。将这个式子设为f(nk,r)$f(nk,r)$，可以得到状态转移方程：f(i,j)=f(i−1,(j−1)modk)+f(i−1,j)$f(i,j)=f(i-1,(j-1)\mathrm{mod}k)+f(i-1,j)$。我们可以把f(i,0),f(i,1),...,f(i,k−1)$f(i,0),f(i,1),...,f(i,k-1)$从上到下排列成一个列向量Fi$F_i$，发现很容易可以构造出一个转移矩阵A$A$使得Fi=AFi−1$F_i=A{F}_{i-1}$，所以Fnk=AnkF0$F_{nk}=A^{nk}{F}_0$，F0$F_0$就是一个第一行为1，其余行为0的列向量。那么我们就可以用矩阵快速幂来加速DP了，总复杂度为O(k3logn)$O(k^3\log n)$，不过据说还有更好的方法能达到O(k2logn)$O(k^2\log n)$，甚至O(klogklogn)$O(k\log k\log n)$，不过这不在这篇文章的讨论范围内了（其实是还没学23333）。要注意k=1$k=1$的特殊情况，在这种情况下，矩阵乘法时要将元素再加一次，因为这个WA了挺久。
以下是本人代码：

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>
#include <algorithm>
#define ll long long
using namespace std;
ll n,r,k,p;
struct matrix {ll a[50][50];} M[50],A;

matrix mult(matrix A,matrix B)
{
  matrix S;
  memset(S.a,0,sizeof(S.a));
  for(int i=0;i<k;i++)
    for(int j=0;j<k;j++)
      for(int l=0;l<k;l++)
        S.a[i][j]=(S.a[i][j]+A.a[i][l]*B.a[l][j])%p;
  if (k==1) S.a[0][0]=(S.a[0][0]+A.a[0][0]*B.a[0][0])%p;
  return S;
}

matrix power(ll x)
{
  matrix S;
  memset(S.a,0,sizeof(S.a));
  for(int i=0;i<k;i++) S.a[i][i]=1;
  int j=0;
  while(x)
  {
    if (x&1) S=mult(S,M[j]);
    x>>=1;j++;
  }
  return S;
}

int main()
{
  scanf("%lld%lld%lld%lld",&n,&p,&k,&r);
  memset(M[0].a,0,sizeof(M[0].a));
  for(int i=0;i<k;i++)
  {
    M[0].a[i][(i-1+k)%k]=1;
    M[0].a[i][i]=1;
  }

  for(int i=1;i<=45;i++) M[i]=mult(M[i-1],M[i-1]);
  A=power(n*k);
  printf("%lld",A.a[r][0]);

  return 0;
}