小圆前辈的素数——FFT_素数长度数列的fft-优快云博客

本文链接：https://blog.youkuaiyun.com/yezi_coder/article/details/119520298

题目大意

给你两组数第一组 $n$ 个，第二组 $m$ 个，在第一组数中选一个数，第二组数中选一个数，问你这两个数之和是素数有多少种选法

解题思路

暴力的方法：

遍历第一组，然后遍历第二组，进行判断，复杂度 $O (n * m)$ 对于数据量大于 $10^5$ 的问题显然无法解决。

FFT：

我们用两个桶来记录第一二组数，然后对两个桶进行FFT
$A(x) = x^{a_1} + x^{a_2} + x^{a_3} ...... x^{a_n} \\ B(x) = x^{b_1} + x^{b_2} + x^{b_3} ...... x^{b_m}$
则对于 $C = A + B$ ，转化为 $C (x) = A (x) * B (x)$
$=\sum _{\forall (i,j) \in [1,1] \times [n,m]} x^{a_i + b_j}$

因此，我们只要将一个多项式中的 $x^{a_n} 和 x^{b_m}$ 的系数赋值为 1 ，其他赋值为 0。进行一遍FFT后，检验 $x^k$ 的系数，若 $x^k$ 的系数大于等于 1 ，则 $\in C$ 。
这样我们就可以在 $O (n l o g n)$ 的时间内算出两组数中任意两个数之和出现的次数。
然后我们筛一遍素数，就可以得到答案了。

Code

#include <bits/stdc++.h>
#define ll long long
#define qc ios::sync_with_stdio(false); cin.tie(0);cout.tie(0)
#define fi first
#define se second
#define PII pair<int, int>
#define PLL pair<ll, ll>
#define pb push_back
using namespace std;
const int MAXN = 3e5 + 7;
const int inf = 0x3f3f3f3f;
const ll INF = 0x3f3f3f3f3f3f3f3f;
inline int read()
{
    int x=0,f=1;char ch=getchar();
    while (!isdigit(ch)){if (ch=='-') f=-1;ch=getchar();}
    while (isdigit(ch)){x=x*10+ch-48;ch=getchar();}
    return x*f;
}
bool cheak[MAXN];//素数为false 
int prime[MAXN];//存第i个素数的值 
int cnt;//素数个数
void Prime()
{
	for(int i = 2; i < MAXN; i++)
	{
		if(!cheak[i]){
			prime[++cnt] = i;
		}
			for(int j=1;j<=cnt&&i*prime[j]<MAXN;j++){
				cheak[i*prime[j]] =true;
				if(i*prime[j] == 0)
				break;
			}
	}
}
const double PI = acos(-1.0);
struct Complex
{
    double x, y;
    Complex operator+(const Complex &W) const
    {
        return {x + W.x, y + W.y};
    }
    Complex operator-(const Complex &W) const
    {
        return {x - W.x, y - W.y};
    }
    Complex operator*(const Complex &W) const
    {
        return {x * W.x - y * W.y, x * W.y + y * W.x};
    }
};
int n, m;
Complex a[MAXN], b[MAXN];
int R[MAXN];
int tot, bit;
void inif(int n)
{
    tot = 1, bit = 0;
    while (tot <= n)
        tot <<= 1, ++bit;
    for (int i = 0; i <= tot; ++i)
        R[i] = (R[i >> 1] >> 1) | ((i & 1) << (bit - 1));
}
void FFT(Complex f[], int total, int type, int n, int m)
{
    for (int i = 0; i < total; ++i)
        if (i < R[i])
            swap(f[i], f[R[i]]);
    for (int tot = 2; tot <= total; tot <<= 1)
    {
        Complex w1 = {cos(2 * PI / tot), type * sin(2 * PI / tot)};
        for (int pos = 0; pos < total; pos += tot)
        {
            Complex w = {1, 0};
            for (int i = pos; i < pos + tot / 2; ++i, w = w * w1)
            {
                Complex x = f[i];
                Complex y = w * f[i + tot / 2];
                f[i] = x + y;
                f[i + tot / 2] = x - y;
            }
        }
    }
    if (type == -1)
    {
        for (int i = 0; i <= n + m; ++i)
            f[i].x = (ll)(f[i].x / tot + 0.5);
    }
}
void solve(){
    cin >> n >> m;
    
    int maxx1 = 0, maxx2 = 0;
    for(int i = 1; i <= n; i++){
    	int x;
    	cin >> x;
    	a[x].x++;
    	maxx1 = max(maxx1, x);
    } 
    for(int i = 1; i <= m; i++){
    	int x;
    	cin >> x;
    	b[x].x++;
    	maxx2 = max(maxx2, x);
    }
    inif(maxx2 + maxx1);
    FFT(a, tot, 1, maxx1, maxx2), FFT(b, tot, 1, maxx1, maxx2);
    for (int i = 0; i < tot; ++i)
        a[i] = a[i] * b[i];
    FFT(a, tot, -1, maxx1, maxx2);
	ll ans = 0;
	for(int i = 1; i <= cnt && prime[i] <= maxx1 + maxx2; i++) {
		ans += a[prime[i]].x;
	}
	cout << ans << endl; 
	for(int i = 0; i <= tot; i++)
    	a[i].x = a[i].y = b[i].x = b[i].y = 0;
}

int main()
{
	#ifdef ONLINE_JUDGE
    #else
       freopen("in.txt", "r", stdin);
       freopen("out.txt", "w", stdout);
    #endif
       Prime();
	qc;
    int T;
    cin >> T;
    while(T--){

        solve();
    }
	return 0;
}