Please subscribe to the official Codeforces channel in Telegram via the link ×

Introducing the imaginary cyclic convolution, speeding up convolution by a factor of 2

Revision en1, by pajenegod, 2022-09-16 14:17:53

I recently had a very interesting idea for how to greatly speed up convolution (a.k.a. polynomial multiplication).

def convolution(A, B):
  C = len(A) + len(B) - 1
  for i in range(len(A)):
    for j in range(len(B)):
      C[i + j] += A[i] * B[j]
  return C

The standard technique for how to do convolution fast is to make use of cyclic convolution (polynomial mult mod $$$x^n - 1$$$).

def cyclic_convolution(A, B):
  n = len(A)
  C = n
  for i in range(n):
    for j in range(n):
      C[(i + j) % n] += A[i] * B[j]
  return C

Cyclic convolution can be calculated in $$$O(n \log n)$$$ using FFT, which is really fast. The issue here is that in order to do convolution using cyclic convolution, we need to pad with a lot of 0s to not be affected by the wrap around. All this 0-padding feels very inefficient.

So here is my idea. What if we do polynomial multiplication mod $$$x^n - i$$$ instead of mod $$$x^n - 1$$$? Then when we get wrap around, it will be multiplied by the imaginary unit, so it wont interfere with the real part! I call this the imaginary cyclic convolution.

def imaginary_cyclic_convolution(A, B):
  n = len(A)
  for i in range(n):
    for j in range(n):
      C[(i + j) % n] += (1 if i + j < n else 1i) * A[i] * B[j]
  return C

Imaginary cyclic convolution is the perfect algorithm to use for implementing convolution. Using it, we no longer need to do copious amount of 0 padding, since the imaginary unit will take care of the wrap around. In fact, the size (the value of $$$n$$$) required is exactly half of what we would need if we had used cyclic convolution.

One question remains, how do we implement imaginary cyclic convolution efficiently?

The trick is rather simple. Let $$$\omega = i^{\frac{1}{n}}$$$. Now note that if $$$C(\omega x) = A(\omega x) B(\omega x) \mod x^n - 1$$$ then $$$C(x) = A(x) B(x) \mod x^n - i$$$. So here is the algorithm

def imaginary_cyclic_convolution(A, B):
  n = len(A)
  omega = (1i)**(1/n)
  # Weight A and B
  A = [A[i] * w**i for i in range(n)]
  B = [B[i] * w**i for i in range(n)]

  C = circular_convolution(A, B)
  # Unweight C
  C = [C[i] / w**i for i in range(n)]
  return C
Tags fft, convolution, constant factor, tricks


  Rev. Lang. By When Δ Comment
en9 English pajenegod 2022-09-16 16:34:17 0 (published)
en8 English pajenegod 2022-09-16 16:33:59 7
en7 English pajenegod 2022-09-16 15:12:46 2 Tiny change: '< n else 1i) * A[i] *' -> '< n else 1j) * A[i] *'
en6 English pajenegod 2022-09-16 14:57:46 126
en5 English pajenegod 2022-09-16 14:36:07 32
en4 English pajenegod 2022-09-16 14:32:50 8 Tiny change: 'n\n C = circular_convoluti' -> 'n\n C = cyclic_convoluti'
en3 English pajenegod 2022-09-16 14:29:11 105
en2 English pajenegod 2022-09-16 14:23:27 6 Tiny change: 'len(A)\n omega = (1i)**(' -> 'len(A)\n w = (1i)**('
en1 English pajenegod 2022-09-16 14:17:53 2316 Initial revision (saved to drafts)