题解：P17217 [ICPC 2017 Nanning R] Five Dimensional Discrete Fourier Transform

· · 题解

题意简述

给定带异或系数与相位因子的五维复数组。求其离散傅里叶变换中，所有频域元素实部绝对值的归一化总和。

解题思路

直接计算五维离散傅里叶变换（Discrete Fourier Transform，DFT）需要枚举两组五维下标。关键是把异或值拆成常数项和 4 个可分离项。

定义第 b 位的符号函数：

\beta_b(x)=(-1)^{\lfloor x/2^b\rfloor\bmod 2}

异或结果的第 b 位为奇数个对应位取 1。因此：

\operatorname{xor}(i_1,\dots,i_5)=\sum_{b=0}^3 2^{b-1}\left(1-\prod_{d=1}^5\beta_b(i_d)\right)

右侧共有一个常数项和 4 个乘积项。原数组中的复指数也能按维度相乘。DFT 具有线性性，因此每个乘积项的五维变换，等于 5 个一维变换的乘积。

令 \sigma_d=(-1)^{d+1}，并定义：

B_{d,b}(j)=\sum_{i=0}^{n_d-1}\beta_b(i)e^{\mathrm{i}\sigma_d\alpha i-2\pi\mathrm{i}ij/n_d}

再令 B_{d,-1} 中的 \beta_{-1}(i)=1。所求频域元素可写成：

A(j_1,\dots,j_5)=\frac{15}{2}\prod_{d=1}^5B_{d,-1}(j_d)-\sum_{b=0}^3 2^{b-1}\prod_{d=1}^5B_{d,b}(j_d)

每个一维变换的长度至多为 10，可以直接预处理。代码把前两维和后三维的乘积分别存下。枚举频域下标时，仅需对 5 个分离项求和。

所有下标都是整数，e^{\mathrm{i}\alpha i} 只依赖于 \alpha\bmod 2\pi。代码先用 long double 约化 \alpha，再计算三角函数，避免大相位降低精度。

设 N=\prod_{d=1}^5n_d。总复杂度为 O(N+\sum n_d^2)，空间复杂度为 O(n_1n_2+n_3n_4n_5)，其中常数均包含 5 个分离项。

参考代码

#include <bits/stdc++.h>
using namespace std;

using ld=long double;
using cd=complex<double>;
const int K=5;
const int N=15;
const int L=105;
const int R=1005;
const ld lpi=acosl(-1);
const double pi=acos(-1);
const double cf[K]={7.5,-0.5,-1,-2,-4};
const int sgn[K]={1,-1,1,-1,1};
cd f[K][K][N],lf[L][K],rf[R][K];
void calc(int d,int n,double a)
{
    for(int i=0;i<K;i++)
    {
        for(int j=0;j<n;j++)
        {
            f[d][i][j]=0;
            for(int k=0;k<n;k++)
            {
                double w=i&&((k>>(i-1))&1)?-1:1,x=(sgn[d]*a-2*pi*j/n)*k;
                f[d][i][j]+=w*cd(cos(x),sin(x));
            }
        }
    }
}
int main()
{
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int T;
    cin>>T;
    cout<<fixed<<setprecision(6);
    while(T--)
    {
        int n[K];
        for(int i=0;i<K;i++)cin>>n[i];
        ld a;
        cin>>a;
        double ang=remainderl(a,2*lpi);
        for(int i=0;i<K;i++)calc(i,n[i],ang);
        int lc=n[0]*n[1];
        for(int i=0;i<lc;i++)
        {
            int x=i/n[1],y=i%n[1];
            for(int j=0;j<K;j++)lf[i][j]=cf[j]*f[0][j][x]*f[1][j][y];
        }
        int rc=n[2]*n[3]*n[4];
        for(int i=0;i<rc;i++)
        {
            int x=i/(n[3]*n[4]),y=i/n[4]%n[3],z=i%n[4];
            for(int j=0;j<K;j++)rf[i][j]=f[2][j][x]*f[3][j][y]*f[4][j][z];
        }
        double ans=0;
        for(int i=0;i<lc;i++)
        {
            for(int j=0;j<rc;j++)
            {
                cd cur=0;
                for(int k=0;k<K;k++)cur+=lf[i][k]*rf[j][k];
                ans+=abs(cur.real());
            }
        }
        int siz=1;
        for(int i=0;i<K;i++)siz*=n[i];
        cout<<ans/(siz*sqrt(1.0*siz))<<'\n';
    }
    return 0;
}