20200412 T2 带mod的FFT

题目描述：

给定数组$a_{N-1,M-1}$aN−1,M−1​和$b_{N-1}$bN−1​，定义：
$c_i=\sum_{j=0}^{N-1}a_{j,b_{i*j\bmod N}}$ci​=j=0∑N−1​aj,bi∗jmodN​​
求第$K$K大的$c_i$ci​。
$N\le249973,b_i<M\le4,a_i\in[0,1024]$N≤249973,bi​<M≤4,ai​∈[0,1024]
保证$N$N为质数。（原题目并没有强调这一点，而是给出了N的表格，坑）

题目分析：

首先特判$i,j=0$i,j=0的情况，然后考虑$ij\neq0$ij​=0的贡献。
注意到 $N$N 都是质数，使用原根把 $i*j\bmod N$i∗jmodN 转化为 $i+j\bmod N-1$i+jmodN−1
设$C_i=c_{g^i},A_{i,j}=a_{g^i,j},B_i=b_{g^i}$Ci​=cgi​,Ai,j​=agi,j​,Bi​=bgi​（其中$C$C是除去$i,j=0$i,j=0时的贡献）
那么有
$C_i=\sum_{j=0}^{N-2}A_{j,B_{i+j\bmod N-1}}$Ci​=j=0∑N−2​Aj,Bi+jmodN−1​​枚举$B$B的取值，转换成艾弗森表达式：
$C_i=\sum_{k=0}^{M-1}\sum_{j=0}^{N-2}A_{j,k}[B_{i+j\bmod N-1}=k]$Ci​=k=0∑M−1​j=0∑N−2​Aj,k​[Bi+jmodN−1​=k]将$B$B翻转后变成卷积形式：
$C_i=\sum_{k=0}^{M-1}\sum_{j=0}^{N-2}A_{j,k}[B_{N-2-(i+j\bmod N-1)}=k]$Ci​=k=0∑M−1​j=0∑N−2​Aj,k​[BN−2−(i+jmodN−1)​=k]，这里的$i+j$i+j虽然带mod，仍然可以FFT，对于一个固定的 $i$i，产生贡献的 $j$j 可以这么看（下面的B表示如果那一位等于k则为1否则为0，A省去了第二维）：
$\sum_{j=0}^{N-2-i}A_{j}B_{N-2-i-j}+\sum_{j=N-1-i}^{N-2}A_jB_{N-2-i-j+N-1}$j=0∑N−2−i​Aj​BN−2−i−j​+j=N−1−i∑N−2​Aj​BN−2−i−j+N−1​前两个多项式相乘时，前者会加到第$N-2-i$N−2−i项，后者会加到$N-2-i+N-1$N−2−i+N−1项，并且可以发现这样不会漏算重算。所以$C_i=(A*B)[N-2-i]+(A*B)[N-2-i+N-1]$Ci​=(A∗B)[N−2−i]+(A∗B)[N−2−i+N−1]。（中括号表示第几项的系数）

Code：

#include<bits/stdc++.h>
#define maxn 524305
using namespace std;
char cb[1<<20],*cs,*ct;
#define getc() (cs==ct&&(ct=(cs=cb)+fread(cb,1,1<<20,stdin),cs==ct)?0:*cs++)
inline void read(int &a){
	char c;while(!isdigit(c=getc()));
	for(a=c-'0';isdigit(c=getc());a=a*10+c-'0');
}
const double Pi = acos(-1);
struct cp{
	double r,i; cp(){}
	cp(double r,double i):r(r),i(i){}
	cp operator + (const cp &t)const{return cp(r+t.r,i+t.i);}
	cp operator - (const cp &t)const{return cp(r-t.r,i-t.i);}
	cp operator * (const cp &t)const{return cp(r*t.r-i*t.i,r*t.i+i*t.r);}
}A[maxn],B[maxn],C[maxn],w[maxn];
int r[maxn];
void FFT(cp *a,int len,int flg){
	for(int i=0;i<len;i++) if(i<r[i]) swap(a[i],a[r[i]]);
	for(int i=2,l=1;i<=len;l=i,i<<=1){
		cp wn = cp(cos(2*Pi/i),sin(2*Pi*flg/i));
		for(int k=l-2;k>=0;k-=2) w[k+1]=(w[k]=w[k>>1])*wn;
		for(int j=0;j<len;j+=i)
			for(int k=j;k<j+l;k++){
				cp u=a[k],v=w[k-j]*a[k+l];
				a[k]=u+v,a[k+l]=u-v;
			}
	}
	if(flg==-1) for(int i=0;i<len;i++) a[i].r/=len;
}
int n,m,K,a[maxn][4],b[maxn],c[maxn],pw[maxn],G;
int p[maxn],cnt;
inline int Pow(int a,int b,const int mod){
	int s=1; for(;b;b>>=1,a=1ll*a*a%mod) if(b&1) s=1ll*s*a%mod;
	return s;
}
int prime_root(const int N){
	int x=N-1;
	for(int i=2;i*i<=x;i++) if(x%i==0) {p[++cnt]=i;while(x%i==0) x/=i;}
	if(x>1) p[++cnt]=x;
	for(int g=2;;g++){
		bool flg=1;
		for(int i=1;i<=cnt;i++) if(Pow(g,(N-1)/p[i],N)==1) {flg=0;break;}
		if(flg) return g;
	}
}
int main()
{
	freopen("b.in","r",stdin);
	freopen("b.out","w",stdout);
	read(n),read(m),read(K),G=prime_root(n);
	for(int i=(pw[0]=1);i<n;i++) pw[i]=1ll*pw[i-1]*G%n;
	for(int i=0;i<m;i++) for(int j=0;j<n;j++) read(a[j][i]);
	for(int i=0;i<n;i++) read(b[i]);
	for(int i=0;i<n;i++) c[0]+=a[i][b[0]];
	for(int i=1;i<n;i++) c[i]+=a[0][b[0]];
	int len=1; while(len<=(n-2)<<1) len<<=1; w[0]=cp(1,0);
	for(int i=0;i<len;i++) r[i]=r[i>>1]>>1|(i&1?len>>1:0);
	for(int k=0;k<m;k++){
		for(int i=0;i<len;i++) A[i]=B[i]=cp(0,0);
		for(int i=0;i<n-1;i++) A[i].r=a[pw[i]][k],B[i].r=(b[pw[i]]==k);
		reverse(B,B+n-1);
		FFT(A,len,1),FFT(B,len,1);
		for(int i=0;i<len;i++) C[i]=A[i]*B[i];
		FFT(C,len,-1);
		for(int i=0;i<n-1;i++) c[pw[i]]+=round(C[n-2-i].r)+round(C[n-2-i+n-1].r);
	}
	sort(c,c+n);
	printf("%d\n",c[n-K]);
}

拓展：

求$C_i=\sum_{j=0}^NA_jB_{(i+j)\bmod M}$Ci​=j=0∑N​Aj​B(i+j)modM​
同样，翻转 $B$B，然后分开来看，大致对应下图：

即$C_i=\sum_{k=1}^{kM-1-i<N+M}(A*revB)[kM-1-i]$Ci​=k=1∑kM−1−i<N+M​(A∗revB)[kM−1−i]