#4736. Tritwise Mex

题目描述

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

题解

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

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

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

所以发现我们可以把原式的操作看成一次卷积，设 $A_0(x)=sum_ia[0 imes 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;
}