コード例 #1
0
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
コード例 #2
0
ファイル: convolution.py プロジェクト: KonstantinTogoi/sympy
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