Skip to content

2021-10-04

백준 1067 - 이동

FFT를 쓰는 알고리즘이라서 시도하고 있습니다.

Fourier series => Fourier Transform => DFT => FFT

순으로 이해

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
import sys
import math
sys.stdin = open("1067_input.txt", "rt")

N = int(input())
As = list(map(int, input().split()))
Bs = list(map(int, input().split()))

def DFT(N, arr):
    ret = [0]*N
    for k in range(N):
        for n in range(N):
            tmp_val = -2*math.pi*k*n/N
            ret[k] += arr[n]*(math.cos(tmp_val) + 1j*math.sin(tmp_val))
    return ret

def IDFT(N, arr):
    ret = [0]*N
    for k in range(N):
        for n in range(N):
            tmp_val = 2*math.pi*k*n/N
            ret[k] += arr[n]*(math.cos(tmp_val) + 1j*math.sin(tmp_val))
        ret[k] /= N
    return ret

def magnitude(arr):
    return list(map(lambda x: (x.real**2 + x.imag**2)**0.5, arr))

As = DFT(N, As)
Bs = DFT(N, Bs)
Cs = []
for i in range(N):
    Cs.append(As[i]*Bs[i])
Cs = IDFT(N, Cs)

list(map(lambda x: x.real, Cs))