def convolution_ntt(a, b, prime):
"""
Performs linear convolution using Number Theoretic Transform.
Parameters
==========
a, b : iterables
The sequences for which convolution is performed.
prime : Integer
Prime modulus of the form (m*2**k + 1) to be used for performing
NTT on the sequence.
Examples
========
>>> from sympy.discrete.convolution import convolution_ntt
>>> convolution_ntt([2, 3], [4, 5], prime=19*2**10 + 1)
[8, 22, 15]
>>> convolution_ntt([2, 5], [6, 7, 3], prime=19*2**10 + 1)
[12, 44, 41, 15]
>>> convolution_ntt([333, 555], [222, 666], prime=19*2**10 + 1)
[15555, 14219, 19404]
References
==========
.. [1] https://en.wikipedia.org/wiki/Convolution_theorem
.. [2] https://en.wikipedia.org/wiki/Discrete_Fourier_transform_(general)
"""
a, b, p = a[:], b[:], as_int(prime)
n = m = len(a) + len(b) - 1 # convolution size
if n > 0 and n & (n - 1): # not a power of 2
n = 2**n.bit_length()
# padding with zeros
a += [0] * (n - len(a))
b += [0] * (n - len(b))
a, b = ntt(a, prime), ntt(b, prime)
for i in range(n):
a[i] = a[i] * b[i] % p
a = intt(a, p)[:m]
return a
def convolution_ntt(a, b, prime):
"""
Performs linear convolution using Number Theoretic Transform.
Parameters
==========
a, b : iterables
The sequences for which convolution is performed.
prime : Integer
Prime modulus of the form (m*2**k + 1) to be used for performing
NTT on the sequence.
Examples
========
>>> from sympy.discrete.convolution import convolution_ntt
>>> convolution_ntt([2, 3], [4, 5], prime=19*2**10 + 1)
[8, 22, 15]
>>> convolution_ntt([2, 5], [6, 7, 3], prime=19*2**10 + 1)
[12, 44, 41, 15]
>>> convolution_ntt([333, 555], [222, 666], prime=19*2**10 + 1)
[15555, 14219, 19404]
References
==========
.. [1] https://en.wikipedia.org/wiki/Convolution_theorem
.. [2] https://en.wikipedia.org/wiki/Discrete_Fourier_transform_(general)
"""
a, b, p = a[:], b[:], as_int(prime)
n = m = len(a) + len(b) - 1 # convolution size
if n > 0 and n&(n - 1): # not a power of 2
n = 2**n.bit_length()
# padding with zeros
a += [0]*(n - len(a))
b += [0]*(n - len(b))
a, b = ntt(a, p), ntt(b, p)
a = [x*y % p for x, y in zip(a, b)]
a = intt(a, p)[:m]
return a