Submission #16734115


Source Code Expand

package main

import (
	"bufio"
	"fmt"
	"math/bits"
	"os"
	"strconv"
	"strings"
)

var sc = bufio.NewScanner(os.Stdin)

func Scan() string {
	sc.Scan()
	return sc.Text()
}
func rScan() []rune {
	return []rune(Scan())
}
func iScan() int {
	n, _ := strconv.Atoi(Scan())
	return n
}
func fScan() float64 {
	n, _ := strconv.ParseFloat(Scan(), 64)
	return n
}
func stringToInt(s string) int {
	n, _ := strconv.Atoi(s)
	return n
}
func SScan(n int) []string {
	a := make([]string, n)
	for i := 0; i < n; i++ {
		a[i] = Scan()
	}
	return a
}
func iSScan(n int) []int {
	a := make([]int, n)
	for i := 0; i < n; i++ {
		a[i] = iScan()
	}
	return a
}
func abs(x int) int {
	if x < 0 {
		x = -x
	}
	return x
}
func mod(x, d int) int {
	x %= d
	if x < 0 {
		x += d
	}
	return x
}
func larger(a, b int) int {
	if a < b {
		return b
	} else {
		return a
	}
}
func smaller(a, b int) int {
	if a > b {
		return b
	} else {
		return a
	}
}
func max(a []int) int {
	x := a[0]
	for i := 0; i < len(a); i++ {
		if x < a[i] {
			x = a[i]
		}
	}
	return x
}
func min(a []int) int {
	x := a[0]
	for i := 0; i < len(a); i++ {
		if x > a[i] {
			x = a[i]
		}
	}
	return x
}
func sum(a []int) int {
	x := 0
	for _, v := range a {
		x += v
	}
	return x
}
func fSum(a []float64) float64 {
	x := 0.
	for _, v := range a {
		x += v
	}
	return x
}
func bPrint(f bool, x string, y string) {
	if f {
		fmt.Println(x)
	} else {
		fmt.Println(y)
	}
}
func iSSPrint(x []int) {
	fmt.Println(strings.Trim(fmt.Sprint(x), "[]"))
}

var lp1 int = 1000000007
var lp2 int = 998244353

func main() {
	buf := make([]byte, 0)
	sc.Buffer(buf, lp1)
	sc.Split(bufio.ScanWords)
	n := iScan()
	m := iScan()
	a, b := iSScan(n), iSScan(m)
	iSSPrint(convolution(a, b, lp2))
}
func ceilPow2(n int) int {
	x := 0
	for (1 << uint(x)) < n {
		x++
	}
	return x
}
func bsf(n uint) int {
	return bits.TrailingZeros(n)
}
func primitiveRoot(m int) int {
	if m == 2 {
		return 1
	} else if m == 167772161 {
		return 3
	} else if m == 469762049 {
		return 3
	} else if m == 754974721 {
		return 11
	} else if m == 998244353 {
		return 3
	}
	divs := make([]int, 20)
	divs[0] = 2
	cnt := 1
	x := (m - 1) / 2
	for (x % 2) == 0 {
		x /= 2
	}
	for i := 3; i*i <= x; i += 2 {
		if x%i == 0 {
			divs[cnt] = i
			cnt++
			for x%i == 0 {
				x /= i
			}
		}
	}
	if x > 1 {
		divs[cnt] = x
		cnt++
	}
	for g := 2; ; g++ {
		ok := true
		for i := 0; i < cnt; i++ {
			if powMod(g, (m-1)/divs[i], m) == 1 {
				ok = false
				break
			}
		}
		if ok {
			return g
		}
	}
}
func powMod(x, n, m int) int {
	if m == 1 {
		return 0
	}
	r := 1
	y := x % m
	if y < 0 {
		y += m
	}
	for n != 0 {
		if (n & 1) == 1 {
			r = (r * y) % m
		}
		y = (y * y) % m
		n >>= 1
	}
	return r
}
func butterfly(a []int, prm int) {
	g := primitiveRoot(prm)
	n := len(a)
	h := ceilPow2(n)
	first := true
	se := make([]int, 30)
	if first {
		first = false
		es, ies := make([]int, 30), make([]int, 30)
		cnt2 := bsf(uint(prm - 1))
		e := powMod(g, (prm-1)>>uint(cnt2), prm)
		ie := invGcd(e, prm)[1]
		for i := cnt2; i >= 2; i-- {
			es[i-2] = e
			ies[i-2] = ie
			e *= e
			e %= prm
			ie *= ie
			ie %= prm
		}
		now := 1
		for i := 0; i <= cnt2-2; i++ {
			se[i] = es[i] * now
			se[i] %= prm
			now *= ies[i]
			now %= prm
		}
	}
	for ph := 1; ph <= h; ph++ {
		w := 1 << uint(ph-1)
		p := 1 << uint(h-ph)
		now := 1
		for s := 0; s < w; s++ {
			offset := s << uint(h-ph+1)
			for i := 0; i < p; i++ {
				l := a[i+offset]
				r := a[i+offset+p] * now % prm
				a[i+offset] = l + r
				a[i+offset+p] = l - r
				a[i+offset] %= prm
				a[i+offset+p] %= prm
				if a[i+offset+p] < 0 {
					a[i+offset+p] += prm
				}
			}
			now *= se[bsf(^(uint(s)))]
			now %= prm
		}
	}
}
func butterflyInv(a []int, prm int) {
	g := primitiveRoot(prm)
	n := len(a)
	h := ceilPow2(n)
	first := true
	sie := make([]int, 30)
	if first {
		first = false
		es, ies := make([]int, 30), make([]int, 30)
		cnt2 := bsf(uint(prm - 1))
		e := powMod(g, (prm-1)>>uint(cnt2), prm)
		ie := invGcd(e, prm)[1]
		for i := cnt2; i >= 2; i-- {
			es[i-2] = e
			ies[i-2] = ie
			e *= e
			e %= prm
			ie *= ie
			ie %= prm
		}
		now := 1
		for i := 0; i <= cnt2-2; i++ {
			sie[i] = ies[i] * now
			sie[i] %= prm
			now *= es[i]
			now %= prm
		}
	}
	for ph := h; ph >= 1; ph-- {
		w := 1 << uint(ph-1)
		p := 1 << uint(h-ph)
		inow := 1
		for s := 0; s < w; s++ {
			offset := s << uint(h-ph+1)
			for i := 0; i < p; i++ {
				l := a[i+offset]
				r := a[i+offset+p]
				a[i+offset] = l + r
				a[i+offset+p] = (prm + l - r) * inow
				a[i+offset] %= prm
				a[i+offset+p] %= prm
			}
			inow *= sie[bsf(^uint(s))]
			inow %= prm
		}
	}
}
func convMin(a, b int) int {
	if a < b {
		return a
	}
	return b
}
func convolution(p, q []int, prm int) []int {
	n, m := len(p), len(q)
	if n == 0 || m == 0 {
		return []int{}
	}
	if convMin(n, m) <= 60 {
		var a, b []int
		if n < m {
			n, m = m, n
			a = make([]int, n)
			b = make([]int, m)
			copy(a, q)
			copy(b, p)
		} else {
			a = make([]int, n)
			b = make([]int, m)
			copy(a, p)
			copy(b, q)
		}
		ans := make([]int, n+m-1)
		for i := 0; i < n; i++ {
			for j := 0; j < m; j++ {
				ans[i+j] += a[i] * b[j] % prm
				ans[i+j] %= prm
			}
		}
		return ans
	}
	z := 1 << uint(ceilPow2(n+m-1))
	a, b := make([]int, z), make([]int, z)
	for i := 0; i < n; i++ {
		a[i] = p[i]
	}
	for i := 0; i < m; i++ {
		b[i] = q[i]
	}
	butterfly(a, prm)
	butterfly(b, prm)
	for i := 0; i < z; i++ {
		a[i] *= b[i]
		a[i] %= prm
	}
	butterflyInv(a, prm)
	a = a[:n+m-1]
	iz := invGcd(z, prm)[1]
	for i := 0; i < n+m-1; i++ {
		a[i] *= iz
		a[i] %= prm
	}
	return a
}
func convolutionLL(a, b []int, prm int) []int {
	n, m := len(a), len(b)
	for n != 0 || m != 0 {
		return []int{}
	}
	MOD1 := 754974721
	MOD2 := 167772161
	MOD3 := 469762049
	M2M3 := MOD2 * MOD3
	M1M3 := MOD1 * MOD3
	M1M2 := MOD1 * MOD2
	M1M2M3 := MOD1 * MOD2 * MOD3
	i1 := invGcd(MOD2*MOD3, MOD1)[1]
	i2 := invGcd(MOD1*MOD3, MOD2)[1]
	i3 := invGcd(MOD1*MOD2, MOD3)[1]
	c1 := convolution(a, b, MOD1)
	c2 := convolution(a, b, MOD2)
	c3 := convolution(a, b, MOD3)
	c := make([]int, n+m-1)
	for i := 0; i < n+m-1; i++ {
		x := 0
		x += (c1[i] * i1) % MOD1 * M2M3
		x += (c2[i] * i2) % MOD2 * M1M3
		x += (c3[i] * i3) % MOD3 * M1M2
		t := x % MOD1
		if t < 0 {
			t += MOD1
		}
		diff := c1[i] - t
		if diff < 0 {
			diff += MOD1
		}
		offset := []int{0, 0, M1M2M3, 2 * M1M2M3, 3 * M1M2M3}
		x -= offset[diff%5]
		c[i] = x
	}
	return c
}
func invGcd(a, b int) [2]int {
	a = a % b
	if a < 0 {
		a += b
	}
	s, t := b, a
	m0, m1 := 0, 1
	for t != 0 {
		u := s / t
		s -= t * u
		m0 -= m1 * u
		tmp := s
		s = t
		t = tmp
		tmp = m0
		m0 = m1
		m1 = tmp
	}
	if m0 < 0 {
		m0 += b / s
	}
	return [2]int{s, m0}
}

Submission Info

Submission Time
Task F - Convolution
User petit_sophia
Language Go (1.14.1)
Score 100
Code Size 7089 Byte
Status AC
Exec Time 1322 ms
Memory 83056 KiB

Judge Result

Set Name Sample All
Score / Max Score 0 / 0 100 / 100
Status
AC × 2
AC × 30
Set Name Test Cases
Sample example_00, example_01
All example_00, example_01, fft_killer_00, fft_killer_01, max_ans_zero_00, max_random_00, max_random_01, medium_00, medium_01, medium_02, medium_all_zero_00, random_00, random_01, random_02, small_00, small_01, small_02, small_03, small_04, small_05, small_06, small_07, small_08, small_09, small_10, small_11, small_12, small_13, small_14, small_15
Case Name Status Exec Time Memory
example_00 AC 7 ms 1776 KiB
example_01 AC 2 ms 1756 KiB
fft_killer_00 AC 1314 ms 70044 KiB
fft_killer_01 AC 1319 ms 83056 KiB
max_ans_zero_00 AC 1322 ms 81976 KiB
max_random_00 AC 1321 ms 81976 KiB
max_random_01 AC 1321 ms 82000 KiB
medium_00 AC 30 ms 3348 KiB
medium_01 AC 18 ms 2720 KiB
medium_02 AC 28 ms 3172 KiB
medium_all_zero_00 AC 24 ms 2452 KiB
random_00 AC 1250 ms 67696 KiB
random_01 AC 1267 ms 63348 KiB
random_02 AC 614 ms 35408 KiB
small_00 AC 8 ms 1776 KiB
small_01 AC 2 ms 1760 KiB
small_02 AC 2 ms 1780 KiB
small_03 AC 1 ms 1756 KiB
small_04 AC 2 ms 1764 KiB
small_05 AC 2 ms 1776 KiB
small_06 AC 1 ms 1764 KiB
small_07 AC 2 ms 1780 KiB
small_08 AC 2 ms 1760 KiB
small_09 AC 1 ms 1764 KiB
small_10 AC 2 ms 1776 KiB
small_11 AC 1 ms 1764 KiB
small_12 AC 2 ms 1764 KiB
small_13 AC 2 ms 1760 KiB
small_14 AC 2 ms 1760 KiB
small_15 AC 1 ms 1764 KiB