#4736. Tritwise Mex

题目描述

给出两个长度为 $3^k$ 的数组 $a,b$ ，求数组 $c$ ， $c[i]={\sum }_{me{x}_3(j,k)=i}{a}_j{b}_k$ ，其中 $me{x}_3(i,j)$ 表示将 $i,j$ 化成 $3$ 进制数，每一位都取 $mex$ 后的值。

题解

其实想法挺自然的考场上不知道在干啥。

考虑求出 $c[0\times 3^x+i],c[1\times 3^x+i],c[2\times 3^x+i]$ 的值

$$c[0\times 3^x+i]=\sum _{mex3(j,k)=i}(a[1\times 3^x+j]+a[2\times 3^x+j])(b[1\times 3^x+k]+b[2\times 3^x+k])$$
$$c[1\times 3^x+i]=\sum _{mex3(j,k)=i}(a[0\times 3^x+j]+a[2\times 3^x+j])(b[0\times 3^x+k]+b[2\times 3^x+k])$$$$-a[2\times 3^x+j]b[2\times 3^x+k]$$
$$c[2\times 3^x+i]=\sum _{mex3(j,k)=i}a[0\times 3^x+j]b[1\times 3^x+k]+a[1\times 3^x+j]b[0\times 3^x+k]$$

所以发现我们可以把原式的操作看成一次卷积，设 $A_0(x)={\sum }_{i}a[0\times 3^x+i]$ ，剩下的同理，则
$$C_0(x)=(A_1(x)+A_2(x))(B_1(x)+B_2(x))$$
$$C_1(x)=(A_0(x)+A_2(x))(B_0(x)+B_2(x))-A_2(x)B_2(x)$$
$$C_2(x)=(A_0(x)+A_1(x)+A_2(x))(B_0(x)+B_1(x)+B_2(x))-C_0(x)-C_1(x)$$

于是递归处理即可。

代码

#include <bits/stdc++.h>
#define LL long long
using namespace std;
const int N=6e5+5;int n,m;
LL a[N][15],b[N][15],c[N][15];
void solve(int k,int d){
	if (d==1){
		c[0][d]=(a[1][d]+a[2][d])*(b[1][d]+b[2][d]);
		c[1][d]=(a[0][d]+a[2][d])*(b[0][d]+b[2][d])-a[2][d]*b[2][d];
		c[2][d]=a[0][d]*b[1][d]+a[1][d]*b[0][d];
		return;
	}
	k/=3;
	for (int i=0;i<k;i++)
		a[i][d-1]=a[i+k+k][d],b[i][d-1]=b[i+k+k][d];
	solve(k,d-1);
	for (int i=0;i<k;i++) c[i+k][d]=-c[i][d-1];
	for (int i=0;i<k;i++)
		a[i][d-1]+=a[i+k][d],b[i][d-1]+=b[i+k][d];
	solve(k,d-1);
	for (int i=0;i<k;i++) c[i][d]=c[i][d-1];
	for (int i=0;i<k;i++)
		a[i][d-1]=a[i+k+k][d]+a[i][d],
		b[i][d-1]=b[i+k+k][d]+b[i][d];
	solve(k,d-1);
	for (int i=0;i<k;i++) c[i+k][d]+=c[i][d-1];
	for (int i=0;i<k;i++)
		a[i][d-1]+=a[i+k][d],b[i][d-1]+=b[i+k][d];
	solve(k,d-1);
	for (int i=0;i<k;i++)
		c[i+k+k][d]=c[i][d-1]-c[i+k][d]-c[i][d];
}
int main(){
	cin>>m;n=1;
	for (int i=0;i<m;i++) n*=3;
	for (int i=0;i<n;i++)
		scanf("%lld",&a[i][m]);
	for (int i=0;i<n;i++)
		scanf("%lld",&b[i][m]);solve(n,m);
	for (int i=0;i<n;i++) printf("%lld ",c[i][m]);
	return 0;
}