CompetitivePrograming/fft2.cpp at master · taiseiKMC/CompetitivePrograming · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
/*
FFT Library
written by nagisa
verified AOJ 2560:Point Distance

*/

#include <bits/stdc++.h>

using namespace std;

const double PI=atan2(0,-1);
typedef complex<double> Comp;

const bool is_pow_of2(const int n)
{
	int a = 1;
	while (n >= a)
    {
		if (n == a)
			return true;
		a <<= 1;
	}
	return false;
}

const int least_exp_of2(const int p)
{
    int tmp=1;
    while(p>tmp)
        tmp<<=1;
    return tmp;
}

const vector<Comp> fft(const vector<Comp> &p)
{
    vector<Comp> v(p);
	int n = v.size();
	assert(is_pow_of2(n));

	for (int j = 1, i = 0; j < n - 1; j++)
    {
		for (int k = n >> 1; k >(i ^= k); k >>= 1);
		if (j < i)
            swap(v[i], v[j]);
	}

	for (int m = 2; m <= n; m <<= 1)
    {
		double deg = (-1) * 2 * PI / m;
		Comp r(cos(deg), sin(deg));
		for (int i = 0;i < n;i += m)
        {
			Comp w(1, 0);
			for (int j = i, k = i + m / 2;k < i + m;j++, k++)
            {
				Comp t1 = v[j], t2 = w*v[k];
				v[j] = t1 + t2, v[k] = t1 - t2;
				w *= r;
			}
		}
	}
	return v;
}

const vector<Comp> ifft(const vector<Comp> &p)
{
    vector<Comp> v(p);
	int n = v.size();
	assert(is_pow_of2(n));
	for (int j = 1, i = 0; j < n - 1; j++)
    {
		for (int k = n >> 1; k >(i ^= k); k >>= 1);
		if (j < i)
            swap(v[i], v[j]);
	}
	for (int m = 2; m <= n; m <<= 1)
    {
		double deg = 2 * PI / m;
		Comp r(cos(deg), sin(deg));
		for (int i = 0;i < n;i += m)
        {
			Comp w(1, 0);
			for (int j = i, k = i + m / 2;k < i + m;j++, k++)
            {
				Comp t1 = v[j], t2 = w*v[k];
				v[j] = t1 + t2, v[k] = t1 - t2;
				w *= r;
			}
		}
	}
	for (int i = 0; i < n;++i)
        v[i] *= 1.0 / n;
	return v;
}

const vector<Comp> convolution(const vector<Comp> &P, const vector<Comp> &Q)
{
    vector<Comp> p(least_exp_of2(P.size())*2), q(least_exp_of2(Q.size())*2);
    copy(P.begin(), P.end(), p.begin());
    copy(Q.begin(), Q.end(), q.begin());
	p=fft(p);
    q=fft(q);
    vector<Comp> tmp(p.size());
	for (int i = 0;i < tmp.size();i++)
        tmp[i] = p[i] * q[i];
	return ifft(tmp);
}