Convolution

Share:
  1func convolutionMod(a, b []int, M int) []int {
  2	n1, n2 := len(a), len(b)
  3	n := n1 + n2 - 1
  4	if n1 == 0 || n2 == 0 {
  5		return []int{}
  6	}
  7
  8	z := 1 << ceilPow2(n)
  9	aa, bb := make([]int, z), make([]int, z)
 10	copy(aa, a)
 11	copy(bb, b)
 12	a, b = aa, bb
 13
 14	butterfly(a, M)
 15	butterfly(b, M)
 16	for i := 0; i < z; i++ {
 17		a[i] = a[i] * b[i] % M
 18	}
 19	butterflyInv(a, M)
 20	a = a[:n]
 21	iz := invMod(z, M)
 22	for i := 0; i < n; i++ {
 23		a[i] = a[i] * iz % M
 24		if a[i] < 0 {
 25			a[i] += M
 26		}
 27	}
 28
 29	return a
 30}
 31
 32func convolution(a, b []int) []int {
 33	n1, n2 := len(a), len(b)
 34	n := n1 + n2 - 1
 35	if n1 == 0 || n2 == 0 {
 36		return []int{}
 37	}
 38
 39	MOD1 := 754974721
 40	MOD2 := 167772161
 41	MOD3 := 469762049
 42	M2M3 := MOD2 * MOD3
 43	M1M3 := MOD1 * MOD3
 44	M1M2 := MOD1 * MOD2
 45	M1M2M3 := MOD1 * MOD2 * MOD3
 46
 47	i1 := invMod(M2M3, MOD1)
 48	i2 := invMod(M1M3, MOD2)
 49	i3 := invMod(M1M2, MOD3)
 50
 51	c1 := convolutionMod(a, b, MOD1)
 52	c2 := convolutionMod(a, b, MOD2)
 53	c3 := convolutionMod(a, b, MOD3)
 54
 55	c := make([]int, n)
 56	offset := []int{0, 0, M1M2M3, 2 * M1M2M3, 3 * M1M2M3}
 57
 58	for i := 0; i < n; i++ {
 59		x := 0
 60		x += c1[i] * i1 % MOD1 * M2M3
 61		x += c2[i] * i2 % MOD2 * M1M3
 62		x += c3[i] * i3 % MOD3 * M1M2
 63		diff := c1[i] - x%MOD1
 64		if diff < 0 {
 65			diff += MOD1
 66		}
 67		x -= offset[diff%5]
 68		c[i] = x
 69	}
 70
 71	return c
 72}
 73
 74func primitiveRoot(m int) int {
 75	if m == 2 {
 76		return 1
 77	}
 78	if m == 167772161 || m == 469762049 || m == 998244353 {
 79		return 3
 80	}
 81	if m == 754974721 {
 82		return 11
 83	}
 84	divs := make([]int, 20)
 85	divs[0] = 2
 86	cnt := 1
 87	x := (m - 1) / 2
 88	for x%2 == 0 {
 89		x /= 2
 90	}
 91	for i := 3; i*i <= x; i += 2 {
 92		if x%i == 0 {
 93			divs[cnt] = i
 94			cnt++
 95			for x%i == 0 {
 96				x /= i
 97			}
 98		}
 99	}
100	if x > 1 {
101		divs[cnt] = x
102		cnt++
103	}
104	for g := 2; ; g++ {
105		ok := true
106		for i := 0; i < cnt; i++ {
107			if powMod(g, (m-1)/divs[i], m) == 1 {
108				ok = false
109				break
110			}
111		}
112		if ok {
113			return g
114		}
115	}
116}
117
118func powMod(a, n, M int) int {
119	if M == 1 {
120		return 0
121	}
122	r := 1
123	for n > 0 {
124		if n&1 == 1 {
125			r = r * a % M
126		}
127		a = a * a % M
128		n >>= 1
129	}
130	return r
131}
132
133func invMod(a, M int) int {
134	p, x, u := M, 1, 0
135	for p != 0 {
136		t := a / p
137		a, p = p, a-t*p
138		x, u = u, x-t*u
139	}
140	return x
141}
142
143func ceilPow2(n int) int {
144	x := 0
145	for 1<<x < n {
146		x++
147	}
148	return x
149}
150
151func bsf(n int) int {
152	x := 0
153	for n&(1<<x) == 0 {
154		x++
155	}
156	return x
157}
158
159func butterfly(a []int, M int) {
160	g := primitiveRoot(M)
161	n := len(a)
162	h := ceilPow2(n)
163
164	se := make([]int, 30)
165	es, ies := make([]int, 30), make([]int, 30)
166	cnt2 := bsf(M - 1)
167	e := powMod(g, (M-1)>>cnt2, M)
168	ie := invMod(e, M)
169	for i := cnt2; i >= 2; i-- {
170		es[i-2] = e
171		ies[i-2] = ie
172		e = e * e % M
173		ie = ie * ie % M
174	}
175	now := 1
176	for i := 0; i <= cnt2-2; i++ {
177		se[i] = es[i] * now % M
178		now = now * ies[i] % M
179	}
180	for ph := 1; ph <= h; ph++ {
181		w := 1 << (ph - 1)
182		p := 1 << (h - ph)
183		now := 1
184		for s := 0; s < w; s++ {
185			offset := s << (h - ph + 1)
186			for i := 0; i < p; i++ {
187				l := a[i+offset]
188				r := a[i+offset+p] * now % M
189				a[i+offset] = (l + r) % M
190				a[i+offset+p] = (M + l - r) % M
191			}
192			now = now * se[bsf(^s)] % M
193		}
194	}
195}
196
197func butterflyInv(a []int, M int) {
198	g := primitiveRoot(M)
199	n := len(a)
200	h := ceilPow2(n)
201
202	sie := make([]int, 30)
203	es, ies := make([]int, 30), make([]int, 30)
204	cnt2 := bsf(M - 1)
205	e := powMod(g, (M-1)>>cnt2, M)
206	ie := invMod(e, M)
207	for i := cnt2; i >= 2; i-- {
208		es[i-2] = e
209		ies[i-2] = ie
210		e = e * e % M
211		ie = ie * ie % M
212	}
213	now := 1
214	for i := 0; i <= cnt2-2; i++ {
215		sie[i] = ies[i] * now % M
216		now = now * es[i] % M
217	}
218	for ph := h; ph >= 1; ph-- {
219		w := 1 << (ph - 1)
220		p := 1 << (h - ph)
221		inow := 1
222		for s := 0; s < w; s++ {
223			offset := s << (h - ph + 1)
224			for i := 0; i < p; i++ {
225				l := a[i+offset]
226				r := a[i+offset+p]
227				a[i+offset] = (l + r) % M
228				a[i+offset+p] = (M + l - r) * inow % M
229			}
230			inow = inow * sie[bsf(^s)] % M
231		}
232	}
233}
234