오랜만이다. 저번주 목금은 휴가, 그 다음 토요일엔 제2회 곰곰컵에 참여했고, 일요일엔 그냥 쉰다고 블로그에 잠깐 무심했었다.
이제 다시 포스팅 시작!

BOJ 10531 - Golf Bot 링크
(2022.11.28 기준 P1)
(치팅하면 골프채가 널 가만두지 않아)

문제

골프 로봇은 정확하게 특정 거리만큼 공을 날릴 수 있다. 그 특정 거리는 N개만큼 있다고 했을 때, M개의 거리를 가진 구멍 중 2타 이내에 공을 넣을 수 있는 구멍의 개수 출력

알고리즘

N개의 수 중 2개 이하로 사용하여 만들 수 있는 거리를 구해야 한다.
이를 하나하나 더하면 시간 초과가 나고, FFT를 사용해야 한다.

풀이

N개의 각 수가 차수로 하여금 계수를 1, 나머지는 계수를 0으로 만든 다항식이 있다면, 그 다항식의 제곱을 구하면 가능한 거리를 알 수 있다.

대충 쉽게 설명하자면, 두 다항식의 합성곱을 구하면 가능한 차수엔 계수가 붙게 된다. (복소수 중 실수 부분). 더 쉽게 설명하자면, 두 수의 곱이라고 생각하면 쉽다. 0인 부분도 다른 자리의 곱에 의해 1 이상이 될 수 있으니깐?
이 문제는 가능한 거리가 1로 잡았으니 1 이상이 되면 가능한 거리라고 생각하면 쉽다.
(이게 맞는 설명인가..?)

FFT의 자세한 설명은
여기와 여기를 참고하면 좋다.

코드

import sys; input = sys.stdin.readline
from cmath import exp, pi
from math import ceil, log2

# 거듭제곱을 이용한 FFT
def fft(a, inv):
    n = len(a)
    j = 0
    for i in range(1, n):
        bit = n >> 1
        while j >= bit:
            j -= bit
            bit >>= 1
        j += bit
        if i < j:
            a[i], a[j] = a[j], a[i]

    p = (-2 if inv else 2) * pi
    s = 2
    while s <= n:
        z = exp(complex(0, p / s))
        for i in range(0, n, s):
            w = 1 + 0j
            for j in range(i, i + (s >> 1)):
                tmp = a[j + (s >> 1)] * w
                a[j + (s >> 1)] = a[j] - tmp
                a[j] += tmp
                w *= z
        s <<= 1
    return [x / n for x in a] if inv else a

N = 1 << ceil(log2(400002)) # 합성곱을 구할 두 다항식의 길이 합(200001 * 2)보다 큰, 가장 작은 2의 거듭제곱
A = [0] * N

# 가능한 거리는 1
A[0] = 1
for _ in range(int(input())):
    A[int(input())] = 1

# 2타 이내에 가능한 모든 거리를 구해야 하므로
# 가능한 거리의 리스트인 A와 자기 자신인 A의 합성곱을 구해야 한다.
A = fft(A, False)
for i in range(N):
    A[i] *= A[i]
A = fft(A, True)

# M개의 구멍마다 가능한지 확인
# 만약 가능한 거리라면 실수 부분이 0보다 크다.
answer = 0
for _ in range(int(input())):
    i = int(input())
    if round(A[i].real):
        print(i, round(A[i].real))
        answer += 1
print(answer)

여담

요즘 FFT 공부에 한창 빠져 있다.
개인적으론 진짜 어려운 알고리즘인 것 같지만, 그래도 조금씩 이해하고 있어서 뿌듯하다.