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

· · 题解

题意简述

五维数组的元素由五个下标的异或值与一个可分离的复相位组成。求其五维离散傅里叶变换中，所有频率点实部绝对值之和，再按数组大小归一化。

解题思路

设第 b 个二进制位对应的符号函数为：

\chi_b(x)=(-1)^{\lfloor x/2^b\rfloor\bmod2}

若 x 的第 b 位为 0，则 \chi_b(x)=1；否则为 -1。五个下标异或结果的第 b 位，等于五个对应位的奇偶校验。因此：

\operatorname{bit}_b(i_1\mathbin{\mathrm{xor}}\dots\mathbin{\mathrm{xor}}i_5) =\frac{1-\prod_{k=1}^{5}\chi_b(i_k)}{2}

由于每个 i_k<10，异或值只会用到第 0 至第 3 位。把四个二进制位重新合并：

i_1\mathbin{\mathrm{xor}}\dots\mathbin{\mathrm{xor}}i_5 =\frac{15}{2}-\sum_{b=0}^{3}2^{b-1}\prod_{k=1}^{5}\chi_b(i_k)

原式中的相位同样能够按维度分离。令：

s_k=(-1)^{k+1}

则：

e^{\mathrm{i}(i_1-i_2+i_3-i_4+i_5)\alpha} =\prod_{k=1}^{5}e^{\mathrm{i}s_ki_k\alpha}

异或展开后的常数项与四个乘积项，都是五个一维函数的乘积。离散傅里叶变换具有线性性，而且可分离数组的五维变换等于五个一维变换的乘积。

对每个维度 k，先预处理不带位符号的一维变换：

F_{k,0}(j)=\sum_{i=0}^{n_k-1}e^{\mathrm{i}s_ki\alpha-2\pi\mathrm{i}ij/n_k}

再对四个二进制位预处理：

F_{k,b+1}(j)=\sum_{i=0}^{n_k-1}\chi_b(i)e^{\mathrm{i}s_ki\alpha-2\pi\mathrm{i}ij/n_k}

于是任意频率点 (j_1,\dots,j_5) 的变换值为：

A(j_1,\dots,j_5) =\frac{15}{2}\prod_{k=1}^{5}F_{k,0}(j_k) -\sum_{b=0}^{3}2^{b-1}\prod_{k=1}^{5}F_{k,b+1}(j_k)

代码不为每个频率点重复解码五维坐标，而是逐维扩展一个张量。开始时只有一个元素，其五个分量分别放入系数：

\left(\frac{15}{2},-\frac{1}{2},-1,-2,-4\right)

加入第 k 维时，旧张量的每个元素分别乘上该维所有 F_{k,b}(j)，生成长度扩大 n_k 倍的新张量。处理完五维后，每个元素的五个分量之和就是一个频率点的 A。这样既不保存原五维数组，也不在内层反复解析混合进制下标。

输入的 \alpha 可能很大，但所有下标都是整数。相位 e^{\mathrm{i}\alpha i} 只与 \alpha 对 2\pi 的余数有关。代码先用 remainderl 约化角度，并在全部三角函数和复数运算中使用 long double。

记 P=n_1n_2n_3n_4n_5。一维变换预处理需要 O(\sum n_k^2)，逐维张量扩展的元素总数由 P 主导，所以总时间复杂度为 O(P+\sum n_k^2)。动态保存五个分量，空间复杂度为 O(P)。

正确性证明

符号函数的乘积在五个对应位中有偶数个 1 时为正，奇数个 1 时为负。因此，(1-\prod\chi_b)/2 准确表示异或结果的第 b 位，四位加权求和得到原异或值。

相位指数是五个只依赖单一下标的项之和，所以复指数等于五个一维因子的乘积。异或展开中的每一项也具有相同的可分离形式。对可分离项展开五重 DFT 求和时，可以依次把五个独立求和提出，得到五个一维 DFT 的乘积。因此，频率点公式与题目定义的五维 DFT 完全相同。

逐维张量扩展时，归纳假设旧元素的第 b 个分量等于前若干维对应一维变换的乘积。乘上新维的 F_{k,b}(j) 后，恰好得到多一维的乘积。处理完五维后，五个带系数分量的和就是公式中的 A(j_1,\dots,j_5)，且所有频率坐标组合各生成一次。

程序对每个频率点累加实部绝对值，最后除以 P\sqrt P，与题目要求的归一化因子一致，所以输出正确。

参考代码

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

using ld=long double;
using cd=complex<ld>;
const int K=5;
const int N=10;
const ld pi=acosl(-1);
const ld c[K]={7.5L,-0.5L,-1,-2,-4};
const int sgn[K]={1,-1,1,-1,1};
cd w[K][K][N];
void init(int d,int n,ld a)
{
    for(int i=0;i<K;i++)
    {
        for(int j=0;j<n;j++)
        {
            w[d][i][j]=0;
            for(int k=0;k<n;k++)
            {
                ld s=i&&((k>>(i-1))&1)?-1:1;
                ld x=(sgn[d]*a-2*pi*j/n)*k;
                w[d][i][j]+=s*cd(cosl(x),sinl(x));
            }
        }
    }
}
void mul(array<cd,K> &a,const array<cd,K> &b,int d,int x)
{
    for(int i=0;i<K;i++)a[i]=b[i]*w[d][i][x];
}
int main()
{
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int T;
    cin>>T;
    cout<<fixed<<setprecision(6);
    while(T--)
    {
        int n[K],siz=1;
        for(int i=0;i<K;i++)
        {
            cin>>n[i];
            siz*=n[i];
        }
        ld a;
        cin>>a;
        a=remainderl(a,2*pi);
        for(int i=0;i<K;i++)init(i,n[i],a);
        vector<array<cd,K>> f(1),g;
        for(int i=0;i<K;i++)f[0][i]=c[i];
        for(int i=0;i<K;i++)
        {
            g.resize(f.size()*n[i]);
            for(int j=0;j<f.size();j++)
            {
                for(int k=0;k<n[i];k++)mul(g[j*n[i]+k],f[j],i,k);
            }
            f.swap(g);
            g.clear();
        }
        ld ans=0;
        for(auto &v:f)
        {
            cd z=0;
            for(int i=0;i<K;i++)z+=v[i];
            ans+=fabsl(z.real());
        }
        cout<<ans/(siz*sqrtl(siz))<<'\n';
    }
    return 0;
}