FFT模板(Python)

import numpy as np
import math
def FFT(P):
    n=int(len(P))#最高项是n-1次
    if n==1:#n=1的时候代表最高项次数为0,直接就是常数
        return P
    omega=np.exp(2*np.pi*1j/n)#计算2^n单位根
    Pe,Po=P[::2],P[1::2]#提取奇偶项
    Ye,Yo=FFT(Pe),FFT(Po)#递归求FFT
    y=[0]*n#构造一个长度为n的数组来存储这一回FFT的值
    for k in range(n//2):
        y[k]=Ye[k]+Yo[k]*(omega**k)#相反数的两个函数值可以只由平方函数值推出
        y[k+n//2]=Ye[k]-Yo[k]*(omega**k)
    return y
def IFFT(P):
    n=int(len(P))
    if n==1:
        return P
    omega=np.exp(-2*np.pi*1j/n)
    Pe,Po=P[::2],P[1::2]
    Ye,Yo=IFFT(Pe),IFFT(Po)
    y=[0]*n
    for k in range(n//2):
        y[k]=Ye[k]+Yo[k]*(omega**k)
        y[k+n//2]=Ye[k]-Yo[k]*(omega**k)
    return y
a=str(input().split())#这里获取的输入不知道为什么包含了单引号和方括号
b=str(input().split())
a=a[::-1];b=b[::-1]#倒过来
a=a[2:-2];b=b[2:-2]#去掉多余的单引号和方括号
print(int(a[::-1])*int(b[::-1]))
#这是用python自带的高精度验证FFT大整数的乘法是否正确
length=max(len(a),len(b))
#找出长度的最大值,两个数相乘的位数不会超过最大值的两倍
length=2**(math.ceil(math.log2(length))+1)
#这是找出距离长度最大值的2的次方最近的值,乘以二就是结果位数的上限
A=X=Y=Z=np.zeros(length);B=np.zeros(length)#构建零向量,把字符转到数组去
for i in range(len(a)):
    A[i]=a[i]
for i in range(len(b)):
    B[i]=b[i]
X,Y=FFT(A),FFT(B)#系数表达变成点表达
#print(X)
Z=np.array(X)*np.array(Y)
result=np.array(IFFT(Z))/len(Z)
result=np.array(result)
#print(result)
for i in range(len(result)):
    while int(result[i]+0.5)>=10:#这里+0.5是为了防止精度误差
        result[i]-=10
        result[i+1]+=1;
    pass
#print('then---------->{}'.format(result))
flag=1
for i in range(len(result)):
    #print(result[-i-1].real)
    if abs(result[-i-1].real)<1e-8 and flag==1:#1e8为阈值
        continue
    else:
        flag=0
        print(int(result[-i-1]+0.5),end='')
        pass
    pass
print('\n')
from scipy.fftpack import fft,ifft#这是自带的包
X,Y=fft(A),fft(B)
Z=np.array(X)*np.array(Y)
print(result)
print(ifft(Z))#结果稍微顺序不一样

over

posted @ 2022-02-08 03:41  saionjisekai  阅读(205)  评论(0编辑  收藏  举报