Skip to content

Commit 7840f82

Browse files
Hiromi-AyaseHiromi-Ayase
committed
Convolutionを実装
1 parent 7b8fca7 commit 7840f82

3 files changed

Lines changed: 367 additions & 0 deletions

File tree

‎Convolution/Convolution.java‎

Lines changed: 272 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,272 @@
1+
/**
2+
* Convolution.
3+
*
4+
* @verified https://atcoder.jp/contests/practice2/tasks/practice2_f
5+
*/
6+
public class Convolution {
7+
/**
8+
* Find a primitive root.
9+
*
10+
* @param m A prime number.
11+
* @return Primitive root.
12+
*/
13+
private static int primitiveRoot(int m) {
14+
if (m == 2) return 1;
15+
if (m == 167772161) return 3;
16+
if (m == 469762049) return 3;
17+
if (m == 754974721) return 11;
18+
if (m == 998244353) return 3;
19+
20+
int[] divs = new int[20];
21+
divs[0] = 2;
22+
int cnt = 1;
23+
int x = (m - 1) / 2;
24+
while (x % 2 == 0) x /= 2;
25+
for (int i = 3; (long) (i) * i <= x; i += 2) {
26+
if (x % i == 0) {
27+
divs[cnt++] = i;
28+
while (x % i == 0) {
29+
x /= i;
30+
}
31+
}
32+
}
33+
if (x > 1) {
34+
divs[cnt++] = x;
35+
}
36+
for (int g = 2; ; g++) {
37+
boolean ok = true;
38+
for (int i = 0; i < cnt; i++) {
39+
if (pow(g, (m - 1) / divs[i], m) == 1) {
40+
ok = false;
41+
break;
42+
}
43+
}
44+
if (ok) return g;
45+
}
46+
}
47+
48+
/**
49+
* Power.
50+
*
51+
* @param x Parameter x.
52+
* @param n Parameter n.
53+
* @param m Mod.
54+
* @return n-th power of x mod m.
55+
*/
56+
private static long pow(long x, long n, int m) {
57+
if (m == 1) return 0;
58+
long r = 1;
59+
long y = x % m;
60+
while (n > 0) {
61+
if ((n & 1) != 0) r = (r * y) % m;
62+
y = (y * y) % m;
63+
n >>= 1;
64+
}
65+
return r;
66+
}
67+
68+
/**
69+
* Ceil of power 2.
70+
*
71+
* @param n Value.
72+
* @return Ceil of power 2.
73+
*/
74+
private static int ceil_pow2(int n) {
75+
int x = 0;
76+
while ((1L << x) < n) x++;
77+
return x;
78+
}
79+
80+
/**
81+
* Pre-calculation for NTT.
82+
*
83+
* @param mod Mod.
84+
* @param g Primitive root of mod.
85+
* @return Pre-calculation table.
86+
*/
87+
private static long[] sumE(int mod, int g) {
88+
long[] sum_e = new long[30];
89+
long[] es = new long[30];
90+
long[] ies = new long[30];
91+
int cnt2 = Integer.numberOfTrailingZeros(mod - 1);
92+
long e = pow(g, (mod - 1) >> cnt2, mod);
93+
long ie = pow(e, mod - 2, mod);
94+
for (int i = cnt2; i >= 2; i--) {
95+
es[i - 2] = e;
96+
ies[i - 2] = ie;
97+
e = e * e % mod;
98+
ie = ie * ie % mod;
99+
}
100+
long now = 1;
101+
for (int i = 0; i < cnt2 - 2; i++) {
102+
sum_e[i] = es[i] * now % mod;
103+
now = now * ies[i] % mod;
104+
}
105+
return sum_e;
106+
}
107+
108+
/**
109+
* Pre-calculation for inverse NTT.
110+
*
111+
* @param mod Mod.
112+
* @param g Primitive root of mod.
113+
* @return Pre-calculation table.
114+
*/
115+
private static long[] sumIE(int mod, int g) {
116+
long[] sum_ie = new long[30];
117+
long[] es = new long[30];
118+
long[] ies = new long[30];
119+
120+
int cnt2 = Integer.numberOfTrailingZeros(mod - 1);
121+
long e = pow(g, (mod - 1) >> cnt2, mod);
122+
long ie = pow(e, mod - 2, mod);
123+
for (int i = cnt2; i >= 2; i--) {
124+
es[i - 2] = e;
125+
ies[i - 2] = ie;
126+
e = e * e % mod;
127+
ie = ie * ie % mod;
128+
}
129+
long now = 1;
130+
for (int i = 0; i < cnt2 - 2; i++) {
131+
sum_ie[i] = ies[i] * now % mod;
132+
now = now * es[i] % mod;
133+
}
134+
return sum_ie;
135+
}
136+
137+
/**
138+
* Inverse NTT.
139+
*
140+
* @param a Target array.
141+
* @param sumIE Pre-calculation table.
142+
* @param mod Mod.
143+
*/
144+
private static void butterflyInv(long[] a, long[] sumIE, int mod) {
145+
int n = a.length;
146+
int h = ceil_pow2(n);
147+
148+
for (int ph = h; ph >= 1; ph--) {
149+
int w = 1 << (ph - 1), p = 1 << (h - ph);
150+
long inow = 1;
151+
for (int s = 0; s < w; s++) {
152+
int offset = s << (h - ph + 1);
153+
for (int i = 0; i < p; i++) {
154+
long l = a[i + offset];
155+
long r = a[i + offset + p];
156+
a[i + offset] = (l + r) % mod;
157+
a[i + offset + p] = (mod + l - r) * inow % mod;
158+
}
159+
int x = Integer.numberOfTrailingZeros(~s);
160+
inow = inow * sumIE[x] % mod;
161+
}
162+
}
163+
}
164+
165+
/**
166+
* Inverse NTT.
167+
*
168+
* @param a Target array.
169+
* @param sumE Pre-calculation table.
170+
* @param mod Mod.
171+
*/
172+
private static void butterfly(long[] a, long[] sumE, int mod) {
173+
int n = a.length;
174+
int h = ceil_pow2(n);
175+
176+
for (int ph = 1; ph <= h; ph++) {
177+
int w = 1 << (ph - 1), p = 1 << (h - ph);
178+
long now = 1;
179+
for (int s = 0; s < w; s++) {
180+
int offset = s << (h - ph + 1);
181+
for (int i = 0; i < p; i++) {
182+
long l = a[i + offset];
183+
long r = a[i + offset + p] * now % mod;
184+
a[i + offset] = (l + r) % mod;
185+
a[i + offset + p] = (l - r + mod) % mod;
186+
}
187+
int x = Integer.numberOfTrailingZeros(~s);
188+
now = now * sumE[x] % mod;
189+
}
190+
}
191+
}
192+
193+
/**
194+
* Convolution by ModInt.
195+
*
196+
* @param a Target array 1.
197+
* @param b Target array 2.
198+
* @return Answer.
199+
*/
200+
public static java.util.List<ModIntFactory.ModInt> convolution(
201+
java.util.List<ModIntFactory.ModInt> a,
202+
java.util.List<ModIntFactory.ModInt> b
203+
) {
204+
int mod = a.get(0).mod();
205+
long[] va = a.stream().mapToLong(ModIntFactory.ModInt::value).toArray();
206+
long[] vb = b.stream().mapToLong(ModIntFactory.ModInt::value).toArray();
207+
long[] c = convolution(va, vb, mod);
208+
ModIntFactory factory = new ModIntFactory(mod);
209+
return java.util.Arrays.stream(c).mapToObj(factory::create).collect(java.util.stream.Collectors.toList());
210+
}
211+
212+
/**
213+
* Convolution.
214+
*
215+
* @param a Target array 1.
216+
* @param b Target array 2.
217+
* @param mod Mod(>= 3).
218+
* @return Answer.
219+
*/
220+
public static long[] convolution(long[] a, long[] b, int mod) {
221+
int n = a.length;
222+
int m = b.length;
223+
224+
int z = 1 << ceil_pow2(n + m - 1);
225+
{
226+
long[] na = new long[z];
227+
long[] nb = new long[z];
228+
System.arraycopy(a, 0, na, 0, n);
229+
System.arraycopy(b, 0, nb, 0, m);
230+
a = na;
231+
b = nb;
232+
}
233+
234+
int g = primitiveRoot(mod);
235+
long[] sume = sumE(mod, g);
236+
long[] sumie = sumIE(mod, g);
237+
238+
butterfly(a, sume, mod);
239+
butterfly(b, sume, mod);
240+
for (int i = 0; i < z; i++) {
241+
a[i] = a[i] * b[i] % mod;
242+
}
243+
butterflyInv(a, sumie, mod);
244+
a = java.util.Arrays.copyOf(a, n + m - 1);
245+
246+
long iz = pow(z, mod - 2, mod);
247+
for (int i = 0; i < n + m - 1; i++) a[i] = a[i] * iz % mod;
248+
return a;
249+
}
250+
251+
/**
252+
* Naive convolution. (Complexity is O(N^2)!!)
253+
*
254+
* @param a Target array 1.
255+
* @param b Target array 2.
256+
* @param mod Mod.
257+
* @return Answer.
258+
*/
259+
public static long[] convolutionNaive(long[] a, long[] b, int mod) {
260+
int n = a.length;
261+
int m = b.length;
262+
int k = n + m - 1;
263+
long[] ret = new long[k];
264+
for (int i = 0; i < n; i++) {
265+
for (int j = 0; j < m; j++) {
266+
ret[i + j] += a[i] * b[j] % mod;
267+
ret[i + j] %= mod;
268+
}
269+
}
270+
return ret;
271+
}
272+
}

‎Convolution/ConvolutionTest.java‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
import java.util.Arrays;
2+
import java.util.List;
3+
import java.util.Random;
4+
import java.util.stream.Collectors;
5+
6+
/**
7+
* Convolution example and test.
8+
*/
9+
public class ConvolutionTest {
10+
public static void main(String[] args) {
11+
Random gen = new Random();
12+
int[] primes = {
13+
1053818881,
14+
1051721729,
15+
1045430273,
16+
1012924417,
17+
1007681537,
18+
1004535809,
19+
998244353,
20+
985661441,
21+
976224257,
22+
975175681
23+
};
24+
25+
for (int i = 0; i < 1000; i++) {
26+
int mod = primes[gen.nextInt(primes.length)];
27+
int n = gen.nextInt(500) + 1;
28+
int m = gen.nextInt(500) + 1;
29+
long[] a = new long[n];
30+
long[] b = new long[m];
31+
32+
for (int j = 0; j < n; j++) {
33+
a[j] = gen.nextInt(1000000000);
34+
}
35+
for (int j = 0; j < m; j++) {
36+
b[j] = gen.nextInt(1000000000);
37+
}
38+
ModIntFactory factory = new ModIntFactory(mod);
39+
List<ModIntFactory.ModInt> va = Arrays.stream(a).mapToObj(factory::create).collect(Collectors.toList());
40+
List<ModIntFactory.ModInt> vb = Arrays.stream(b).mapToObj(factory::create).collect(Collectors.toList());
41+
long[] expected = Convolution.convolutionNaive(a, b, mod);
42+
long[] actual = Convolution.convolution(va, vb).stream().mapToLong(ModIntFactory.ModInt::value).toArray();
43+
assertArrayEquals(expected, actual);
44+
}
45+
}
46+
47+
private static void assertArrayEquals(long[] expected, long[] actual) {
48+
if (expected.length != actual.length) throw new RuntimeException();
49+
int n = expected.length;
50+
for (int i = 0; i < n; i++) {
51+
if (expected[i] != actual[i]) {
52+
throw new RuntimeException();
53+
}
54+
}
55+
}
56+
}

‎Convolution/README.md‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
# クラス Convolution
2+
3+
- - -
4+
5+
NTTで剰余Convolutionを計算するクラスです。
6+
7+
## 使い方
8+
9+
1. staticな `Convolution.convolute` を呼び出します.
10+
11+
```java
12+
public static void main(String[] args) {
13+
// 3以上の素数のModを指定します
14+
int mod = 998244353;
15+
long[] a = { 1, 2, 3, 4, 5 };
16+
long[] b = { 6, 7, 8 };
17+
// 畳み込みを計算しまqす。a.length + b.length - 1 の配列が帰ります。
18+
long[] ret = Convolution.convolution(a, b, mod);
19+
System.out.println(Arrays.toString(ret));
20+
}
21+
```
22+
23+
## メソッド
24+
25+
### convolution
26+
27+
```java
28+
public static long[] convolution(long[] a, long[] b, int mod)
29+
public static java.util.List<ModIntFactory.ModInt> convolution(
30+
java.util.List<ModIntFactory.ModInt> a,
31+
java.util.List<ModIntFactory.ModInt> b
32+
)
33+
```
34+
35+
リストまたは配列 `a` と `b` の剰余convolutionを計算します.
36+
計算量: $O(n ¥log(n))$
37+
38+
制約
39+
- `mod` 3以上の素数

0 commit comments

Comments
 (0)