Hi Codeforces!

I have recently come up with a really neat and simple recursive algorithm for multiplying polynomials in $$$O(n \log n)$$$ time. It is so neat and simple that I think it might possibly revolutionize the way that fast polynomial multiplication is taught and implemented.

Prerequisite: Polynomial quotient and remainder, see Wiki article and Stackexchange example.

Task:

Given two polynomials $$$P$$$ and $$$Q$$$ and an integer $$$n$$$, where degree $$$P < n$$$ and degree $$$Q < n$$$. Your task is to calculate the polynomial $$$P(x) \, Q(x) \% (x^n - c)$$$ in $$$O(n \log n)$$$ time. You may assume that $$$n$$$ is a power of two.

Solution:

We can create a divide and conquer algorithm for $$$P(x) \, Q(x) \% (x^n - c)$$$ based on the difference of squares formula. Assuming $$$n$$$ is even, then $$$(x^n - c) = (x^{n/2} - \sqrt{c}) (x^{n/2} + \sqrt{c})$$$. The idea behind the algorithm is to calculate $$$P(x) \, Q(x) \% (x^{n/2} - \sqrt{c})$$$ and $$$P(x) \, Q(x) \% (x^{n/2} + \sqrt{c})$$$ using 2 recursive calls, and then we use them to calculate $$$P(x) \, Q(x) \% (x^n - c)$$$.

So how do we actually calculate $$$P(x) \, Q(x) \% (x^n - c)$$$ using $$$P(x) \, Q(x) \% (x^{n/2} - \sqrt{c})$$$ and $$$P(x) \, Q(x) \% (x^{n/2} + \sqrt{c})$$$?

Well, we can use the following formula:

This formula has a lot of uses. If we substitute $$$A(x)$$$ by $$$P(x) Q(x)$$$, then the formula tells us how to calculate $$$P(x) \, Q(x) \% (x^n - c)$$$ using $$$P(x) \, Q(x) \% (x^{n/2} - \sqrt{c})$$$ and $$$P(x) \, Q(x) \% (x^{n/2} + \sqrt{c})$$$. Additionally, if we substite $$$A(x)$$$ with $$$P(x)$$$, then the formula tells us how to calculate $$$P(x) \% (x^{n/2} - \sqrt{c})$$$ and $$$P(x) \% (x^{n/2} + \sqrt{c})$$$ using $$$P(x) \% (x^n - c)$$$. With this we have entire recipie for implementing a $$$O(n \log n)$$$ divide and conquer algorithm.

"""
Calculates P(x) * Q(x) % (x^n - c) in O(n log n) time

Input:
  n: integer, needs to be power of 2
  P: A list of length n representing a polynomial P(x)
  Q: A list of length n representing a polynomial Q(x)
Output:
  A list of length n representing the polynomial P(x) * Q(x) % (x^n - c)
"""
def fast_polymult(P, Q, n, c):
    assert len(P) == n and len(Q) == n
    
    # Base case
    if n == 1:
        return [P[0] * Q[0]]

    assert n % 2 == 0
    import cmath
    sqrtc = cmath.sqrt(c)

    # Calulate P_minus := P mod (x^(n/2) - sqrt(c))
    #          Q_minus := Q mod (x^(n/2) - sqrt(c))

    P_minus = [p1 + sqrtc * p2 for p1,p2 in zip(P[:n//2], P[n//2:])]
    Q_minus = [q1 + sqrtc * q2 for q1,q2 in zip(Q[:n//2], Q[n//2:])]

    # Calulate P_plus := P mod (x^(n/2) + sqrt(c))
    #          Q_plus := Q mod (x^(n/2) + sqrt(c))

    P_plus = [p1 - sqrtc * p2 for p1,p2 in zip(P[:n//2], P[n//2:])]
    Q_plus = [q1 - sqrtc * q2 for q1,q2 in zip(Q[:n//2], Q[n//2:])]

    # Recursive calculate product
    PQ_minus = polymult_mod_recursive(P_minus, Q_minus, n//2, sqrtc)
    PQ_plus  = polymult_mod_recursive(P_plus,  Q_plus,  n//2, -sqrtc)

    # Calculate result PQ mod (x^n - c)
    PQ = [(m + p)/2       for m,p in zip(PQ_minus, PQ_plus)] +
         [(m + p)/2/sqrtc for m,p in zip(PQ_minus, PQ_plus)]
    
    return PQ