보석 가게

• 고속 푸리에 변환 (Fast Fourier Transform, FFT)을 이용해서 All Possible Sums를 구하는 방식으로 보석 2개를 훔칠때의 가치의 합의 목록을 구할 수 있다. 이 결과에 같은 작업을 또 한번 적용하면 보석 3개를 훔칠때의 가치의 합들을 구할수 있고, 이런식으로 K개를 훔칠때의 합을 계산하면 된다.
• 고속 푸리에 변환 (Fast Fourier Transform, FFT)은 곧 다항식 곱셈이니까, 이것은 다항식의 K제곱이다. 이는 거듭제곱의 빠른 계산 (Exponentiation by squaring) 방법을 이용하면 k제곱을 구하는 것은 O(K)번이 아니라 O(logK)번의 곱셈으로 가능하다. 즉, FFT의 횟수도 O(logK)번. 그러나 항의 길이가 n→2n→4n …으로 계속 증가하므로, FFT 한번에 걸리는 시간도 계속 늘어나서 전체 복잡도는 O(knlog(kn))이 된다.
• 현재 FFT구현이 다항식을 decimal로 변환해서 곱하는 방식으로 되어있는 것처럼, k제곱을 구하는 것도 decimal에 라이브러리의 power함수를 쓰면 바로 되지 않을까 싶지만 거기에는 문제가 있다. 실제로 K제곱을 하면 계수의 최대 크기가 O(n^K)로 엄청나게 커진다. 이만큼의 자리수를 할당해서 decimal을 만드는 것이 불가능하기때문에, 곱셈을 한번 하고 나면 그 결과에서 계수들을 다시 0 또는 1로 바꿔주는 작업을 하는 방식으로 power 함수를 쓰지 말고 구현해야 한다. 한번 곱셈을 할때마다 list를 decimal로 변환하고, decimal끼리 곱하고, 결과를 다시 list로 변환하는 과정을 거쳐야 하기 때문에, 시간이 매우 많이 걸린다. 정말 느렸지만 기준 제한시간이 5초였던 문제였기에, 14000ms에 가까운 시간으로도 통과하기는 했다. (제출하고 채점이 끝날때까지 15분 가까이 걸렸다..)

"""Solution code for "BOJ 13575. 보석 가게".

- Problem link: https://www.acmicpc.net/problem/13575
- Solution link: http://www.teferi.net/ps/problems/boj/13575
"""

from teflib import fft


def main():
    N, K = [int(x) for x in input().split()]  # pylint: disable=unused-variable
    a = [int(x) for x in input().split()]
    jewel = [0] * (max(a) + 1)
    for a_i in a:
        jewel[a_i] = 1

    # Computes answer = jewel ^ K.
    answer = [1]
    for c in reversed(bin(K)):
        if c == '1':
            answer = [0 if x == 0 else 1 for x in fft.multiply(answer, jewel)]
        jewel = [0 if x == 0 else 1 for x in fft.multiply(jewel, jewel)]

    print(' '.join(str(i) for i, x in enumerate(answer) if x == 1))


if __name__ == '__main__':
    main()

• Dependency: teflib.fft.multiply