1. 程式人生 > 實用技巧 >luogu P4717 【模板】快速沃爾什變換 (FWT)

luogu P4717 【模板】快速沃爾什變換 (FWT)

程式碼:

#include<cstdio>
#include<iostream>
#include<cstring>
#include<algorithm>
#define int long long

using namespace std;

const int N=1000000,M=998244353;
int A[N],B[N],tmp_A[N],tmp_B[N],n,inv2;

int ksm(int a,int b)
{
	int res=1;
	while(b)
	{
		if(b&1)
			res=res*a%M;
		b>>=1,a=a*a%M;
	}
	return res;
}

void init()
{
	scanf("%lld",&n);
	n=1<<n;
	for (int i=0;i<n;i++)
		scanf("%lld",&tmp_A[i]);
	for (int i=0;i<n;i++)
		scanf("%lld",&tmp_B[i]);
	inv2=ksm(2,M-2);
}

void fwt_or(int A[],int flag)
{
	for (int l=1;l<n;l<<=1)
		for (int i=0;i<n;i+=l<<1)
			for (int j=0;j<l;j++)
			{
				int u=A[i+j],v=A[i+j+l];
				if(flag==1)
					A[i+j]=u,A[i+j+l]=(u+v)%M;
				else
					A[i+j]=u,A[i+j+l]=(v-u)%M;
			}
}

void FWT_OR()
{
	for (int i=0;i<n;i++)
		A[i]=tmp_A[i],B[i]=tmp_B[i];
	fwt_or(A,1),fwt_or(B,1);
	for (int i=0;i<n;i++)
		A[i]=A[i]*B[i]%M;
	fwt_or(A,-1);
	for (int i=0;i<n;i++)
		printf("%lld ",(A[i]+M)%M);puts("");
}

void fwt_and(int A[],int flag)
{
	for (int l=1;l<n;l<<=1)
		for (int i=0;i<n;i+=l<<1)
			for (int j=0;j<l;j++)
			{
				int u=A[i+j],v=A[i+j+l];
				if(flag==1)
					A[i+j]=(u+v)%M,A[i+j+l]=v;
				else
					A[i+j]=(u-v)%M,A[i+j+l]=v;
			}
}

void FWT_AND()
{
	for (int i=0;i<n;i++)
		A[i]=tmp_A[i],B[i]=tmp_B[i];
	fwt_and(A,1),fwt_and(B,1);
	for (int i=0;i<n;i++)
		A[i]=A[i]*B[i]%M;
	fwt_and(A,-1);
	for (int i=0;i<n;i++)
		printf("%lld ",(A[i]+M)%M);puts("");
}

void fwt_xor(int A[],int flag)
{
	for (int l=1;l<n;l<<=1)
		for (int i=0;i<n;i+=l<<1)
			for (int j=0;j<l;j++)
			{
				int u=A[i+j],v=A[i+j+l];
				if(flag==1)
					A[i+j]=(u+v)%M,A[i+j+l]=(u-v)%M;
				else
					A[i+j]=(u+v)%M*inv2%M,A[i+j+l]=(u-v)%M*inv2%M;
			}
}

void FWT_XOR()
{
	for (int i=0;i<n;i++)
		A[i]=tmp_A[i],B[i]=tmp_B[i];
	fwt_xor(A,1),fwt_xor(B,1);
	for (int i=0;i<n;i++)
		A[i]=A[i]*B[i]%M;
	fwt_xor(A,-1);
	for (int i=0;i<n;i++)
		printf("%lld ",(A[i]+M)%M);puts("");
}

void work()
{
	FWT_OR(),FWT_AND(),FWT_XOR();
}

signed main()
{
//	freopen("P4717.in","r",stdin);
//	freopen("P4717.out","w",stdout);
	init();
	work();
	return 0;
}