fast polynomial multiplication using radix2 fast fourier transform fast polynomial multiplication using radix2 fast fourier transform reference https en wikipedia orgwikicooleye28093tukeyfftalgorithmtheradix2ditcase for polynomials of degree m and n the algorithms has complexity onlogn mlogm the main part of the algorithm is split in two parts 1 dft we compute the discrete fourier transform dft of a and b using a bottomup dynamic approach 2 multiply once we obtain the dft of ab we can similarly invert it to obtain ab the class fft takes two polynomials a and b with complex coefficients as arguments the two polynomials should be represented as a sequence of coefficients starting from the free term thus for instance x 2x3 could be represented as 0 1 0 2 or 0 1 0 2 the constructor adds some zeros at the end so that the polynomials have the same length which is a power of 2 at least the length of their product example create two polynomials as sequences a 0 1 0 2 x2x3 b 2 3 4 0 23x4x2 create an fft object with them x ffta b print product x product 2x 3x2 8x3 4x4 6x5 00j 20j 30j 80j 60j 80j str test printx a 0x0 1x1 2x0 3x2 b 0x2 1x3 2x4 ab 0x00j 1x20j 2x30j 3x80j 4x60j 5x80j input as list remove leading zero coefficients add 0 to make lengths equal a power of 2 a complex root used for the fourier transform the product discrete fourier transform of a and b corner case first half of next step second half of next step update multiply the dfts of a and b and find ab corner case inverse dft first half of next step even positions odd positions update unpack remove leading 0 s overwrite str for print shows a b and ab unit tests for roots of unity fast polynomial multiplication using radix 2 fast fourier transform reference https en wikipedia org wiki cooley e2 80 93tukey_fft_algorithm the_radix 2_dit_case for polynomials of degree m and n the algorithms has complexity o n logn m logm the main part of the algorithm is split in two parts 1 __dft we compute the discrete fourier transform dft of a and b using a bottom up dynamic approach 2 __multiply once we obtain the dft of a b we can similarly invert it to obtain a b the class fft takes two polynomials a and b with complex coefficients as arguments the two polynomials should be represented as a sequence of coefficients starting from the free term thus for instance x 2 x 3 could be represented as 0 1 0 2 or 0 1 0 2 the constructor adds some zeros at the end so that the polynomials have the same length which is a power of 2 at least the length of their product example create two polynomials as sequences a 0 1 0 2 x 2x 3 b 2 3 4 0 2 3x 4x 2 create an fft object with them x fft a b print product x product 2x 3x 2 8x 3 4x 4 6x 5 0 0j 2 0j 3 0j 8 0j 6 0j 8 0j __str__ test print x a 0 x 0 1 x 1 2 x 0 3 x 2 b 0 x 2 1 x 3 2 x 4 a b 0 x 0 0j 1 x 2 0j 2 x 3 0j 3 x 8 0j 4 x 6 0j 5 x 8 0j input as list remove leading zero coefficients add 0 to make lengths equal a power of 2 a complex root used for the fourier transform the product discrete fourier transform of a and b corner case first half of next step second half of next step update multiply the dfts of a and b and find a b corner case inverse dft first half of next step even positions odd positions update unpack remove leading 0 s overwrite __str__ for print shows a b and a b unit tests
import mpmath  
import numpy as np


class FFT:
    

    def __init__(self, poly_a=None, poly_b=None):
        
        self.polyA = list(poly_a or [0])[:]
        self.polyB = list(poly_b or [0])[:]

        
        while self.polyA[-1] == 0:
            self.polyA.pop()
        self.len_A = len(self.polyA)

        while self.polyB[-1] == 0:
            self.polyB.pop()
        self.len_B = len(self.polyB)

        
        self.c_max_length = int(
            2 ** np.ceil(np.log2(len(self.polyA) + len(self.polyB) - 1))
        )

        while len(self.polyA) < self.c_max_length:
            self.polyA.append(0)
        while len(self.polyB) < self.c_max_length:
            self.polyB.append(0)
        
        self.root = complex(mpmath.root(x=1, n=self.c_max_length, k=1))

        
        self.product = self.__multiply()

    
    def __dft(self, which):
        dft = [[x] for x in self.polyA] if which == "A" else [[x] for x in self.polyB]
        
        if len(dft) <= 1:
            return dft[0]
        
        next_ncol = self.c_max_length // 2
        while next_ncol > 0:
            new_dft = [[] for i in range(next_ncol)]
            root = self.root**next_ncol

            
            current_root = 1
            for j in range(self.c_max_length // (next_ncol * 2)):
                for i in range(next_ncol):
                    new_dft[i].append(dft[i][j] + current_root * dft[i + next_ncol][j])
                current_root *= root
            
            current_root = 1
            for j in range(self.c_max_length // (next_ncol * 2)):
                for i in range(next_ncol):
                    new_dft[i].append(dft[i][j] - current_root * dft[i + next_ncol][j])
                current_root *= root
            
            dft = new_dft
            next_ncol = next_ncol // 2
        return dft[0]

    
    def __multiply(self):
        dft_a = self.__dft("A")
        dft_b = self.__dft("B")
        inverce_c = [[dft_a[i] * dft_b[i] for i in range(self.c_max_length)]]
        del dft_a
        del dft_b

        
        if len(inverce_c[0]) <= 1:
            return inverce_c[0]
        
        next_ncol = 2
        while next_ncol <= self.c_max_length:
            new_inverse_c = [[] for i in range(next_ncol)]
            root = self.root ** (next_ncol // 2)
            current_root = 1
            
            for j in range(self.c_max_length // next_ncol):
                for i in range(next_ncol // 2):
                    
                    new_inverse_c[i].append(
                        (
                            inverce_c[i][j]
                            + inverce_c[i][j + self.c_max_length // next_ncol]
                        )
                        / 2
                    )
                    
                    new_inverse_c[i + next_ncol // 2].append(
                        (
                            inverce_c[i][j]
                            - inverce_c[i][j + self.c_max_length // next_ncol]
                        )
                        / (2 * current_root)
                    )
                current_root *= root
            
            inverce_c = new_inverse_c
            next_ncol *= 2
        
        inverce_c = [round(x[0].real, 8) + round(x[0].imag, 8) * 1j for x in inverce_c]

        
        while inverce_c[-1] == 0:
            inverce_c.pop()
        return inverce_c

    
    def __str__(self):
        a = "A = " + " + ".join(
            f"{coef}*x^{i}" for coef, i in enumerate(self.polyA[: self.len_A])
        )
        b = "B = " + " + ".join(
            f"{coef}*x^{i}" for coef, i in enumerate(self.polyB[: self.len_B])
        )
        c = "A*B = " + " + ".join(
            f"{coef}*x^{i}" for coef, i in enumerate(self.product)
        )

        return f"{a}\n{b}\n{c}"



if __name__ == "__main__":
    import doctest

    doctest.testmod()