Algorithm_templates/Convolution.cpp at main · seo-bo/Algorithm_templates · 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
class Convolution //https://github.com/seo-bo/Algorithm_templates/blob/main/Convolution.cpp
{
private:
	vector<complex<double>>v1, v2;
	int len;
	void fft(vector<complex<double>>& v, int pivot)
	{
		for (int i = 1, j = 0; i < len; ++i)
		{
			int bit = len >> 1;
			while (j & bit)
			{
				j ^= bit;
				bit >>= 1;
			}
			j ^= bit;
			if (i < j)
			{
				swap(v[i], v[j]);
			}
		}
		for (int i = 2; i <= len; i <<= 1)
		{
			double angle = 2 * acos(-1.0) * pivot / i;
			complex<double> z(cos(angle), sin(angle));
			for (int j = 0; j < len; j += i)
			{
				complex<double> w(1, 0);
				for (int k = 0; k < i / 2; ++k)
				{
					complex<double>fa = v[j + k];
					complex<double>fb = v[j + k + i / 2] * w;
					v[j + k] = fa + fb;
					v[j + k + i / 2] = fa - fb;
					w *= z;
				}
			}
		}
		if (pivot == -1)
		{
			for (auto& i : v)
			{
				i /= len;
			}
		}
	}
	vector<int> get_conv()
	{
		fft(v1, 1);
		fft(v2, 1);
		for (int i = 0; i < len; ++i)
		{
			v1[i] *= v2[i];
		}
		fft(v1, -1);
		vector<int>ans(len);
		for (int i = 0; i < len; ++i)
		{
			ans[i] = (int)round(v1[i].real());
		}
		return ans;
	}
public:
	Convolution(vector<int>& p1, vector<int>& p2)
	{
		v1 = vector<complex<double>>(p1.begin(), p1.end());
		v2 = vector<complex<double>>(p2.begin(), p2.end());
		int n = 1;
		while (n < (int)v1.size() + (int)v2.size())
		{
			n <<= 1;
		}
		len = n;
		v1.resize(len), v2.resize(len);
	}
	vector<int> solve()
	{
		return get_conv();
	}
};