20200412 T2 帶mod的FFT

題目描述：

給定數組$a_{N-1,M-1}$和$b_{N-1}$，定義：
$c_i=\sum_{j=0}^{N-1}a_{j,b_{i*j\bmod N}}$
求第$K$大的$c_i$。
$N\le249973,b_i<M\le4,a_i\in[0,1024]$
保證$N$爲質數。（原題目並沒有強調這一點，而是給出了N的表格，坑）

題目分析：

首先特判$i,j=0$的情況，然後考慮$ij\neq0$的貢獻。
注意到 $N$ 都是質數，使用原根把 $i*j\bmod N$ 轉化爲 $i+j\bmod N-1$
設$C_i=c_{g^i},A_{i,j}=a_{g^i,j},B_i=b_{g^i}$（其中$C$是除去$i,j=0$時的貢獻）
那麼有
$C_i=\sum_{j=0}^{N-2}A_{j,B_{i+j\bmod N-1}}$枚舉$B$的取值，轉換成艾弗森表達式：
$C_i=\sum_{k=0}^{M-1}\sum_{j=0}^{N-2}A_{j,k}[B_{i+j\bmod N-1}=k]$將$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]$，這裏的$i+j$雖然帶mod，仍然可以FFT，對於一個固定的 $i$，產生貢獻的 $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}$前兩個多項式相乘時，前者會加到第$N-2-i$項，後者會加到$N-2-i+N-1$項，並且可以發現這樣不會漏算重算。所以$C_i=(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}$
同樣，翻轉 $B$，然後分開來看，大致對應下圖：

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