순환 이동 후 최대 점수 구하기 (FFT 활용)

문제 설명

주어진 두 배열 (X)와 (Y)를 순환 이동시켜 각 배열 요소의 곱의 합 (S)를 최대로 만드는 문제이다. 점수 (S)는 다음과 같이 정의된다:

[
S = X[0] \times Y[0] + X[1] \times Y[1] + \dots + X[N-1] \times Y[N-1]
]

여기서, (X)와 (Y)는 각각 여러 번 순환 이동할 수 있으며, (S)의 최댓값을 구해야 한다.

핵심 아이디어

1. 순환 이동의 특징:

   • (X)를 순환 이동하면 (X')가 (X)의 확장된 배열로 표현될 수 있다. 예를 들어, (X = [1, 2, 3])라면 (X' = [1, 2, 3, 1, 2, 3]).
   • (Y)는 반전(reverse) 배열을 만들어 (X)와의 모든 경우를 효율적으로 계산할 수 있다.

2. 빠른 계산을 위한 FFT:

   • (X')와 (Y)의 모든 경우의 점수를 계산하려면 곱셈 합산을 효율적으로 수행해야 한다.
   • 이를 위해 고속 푸리에 변환(FFT)를 사용하여 배열의 곱셈을 빠르게 계산한다.

3. 정확한 최대값 계산:

   • FFT를 이용해 (X')와 (Y)의 곱 결과를 얻고, 이 중 최댓값을 출력한다.

코드

import java.io.*;
import java.util.*;

public class Main {

    static class Complex {
        double real, imag;

        Complex(double real, double imag) {
            this.real = real;
            this.imag = imag;
        }

        Complex add(Complex o) {
            return new Complex(real + o.real, imag + o.imag);
        }

        Complex subtract(Complex o) {
            return new Complex(real - o.real, imag - o.imag);
        }

        Complex multiply(Complex o) {
            return new Complex(real * o.real - imag * o.imag, real * o.imag + imag * o.real);
        }

        Complex divide(double val) {
            return new Complex(real / val, imag / val);
        }
    }

    static void fft(Complex[] a, boolean invert) {
        int n = a.length;
        for (int i = 1, j = 0; i < n; i++) {
            int bit = n >> 1;
            while (j >= bit) {
                j -= bit;
                bit >>= 1;
            }
            j += bit;
            if (i < j) {
                Complex temp = a[i];
                a[i] = a[j];
                a[j] = temp;
            }
        }

        for (int len = 2; len <= n; len <<= 1) {
            double angle = 2 * Math.PI / len * (invert ? -1 : 1);
            Complex wlen = new Complex(Math.cos(angle), Math.sin(angle));
            for (int i = 0; i < n; i += len) {
                Complex w = new Complex(1, 0);
                for (int j = 0; j < len / 2; j++) {
                    Complex u = a[i + j];
                    Complex v = a[i + j + len / 2].multiply(w);
                    a[i + j] = u.add(v);
                    a[i + j + len / 2] = u.subtract(v);
                    w = w.multiply(wlen);
                }
            }
        }

        if (invert) {
            for (int i = 0; i < n; i++) {
                a[i] = a[i].divide(n);
            }
        }
    }

    static long[] multiply(long[] a, long[] b) {
        int n = 1;
        while (n < a.length + b.length) n <<= 1;

        Complex[] fa = new Complex[n];
        Complex[] fb = new Complex[n];

        for (int i = 0; i < n; i++) {
            fa[i] = new Complex(i < a.length ? a[i] : 0, 0);
            fb[i] = new Complex(i < b.length ? b[i] : 0, 0);
        }

        fft(fa, false);
        fft(fb, false);

        for (int i = 0; i < n; i++) {
            fa[i] = fa[i].multiply(fb[i]);
        }

        fft(fa, true);

        long[] result = new long[n];
        for (int i = 0; i < n; i++) {
            result[i] = Math.round(fa[i].real);
        }
        return result;
    }

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int n = Integer.parseInt(br.readLine());

        long[] x = new long[n * 2];
        long[] y = new long[n];

        StringTokenizer st = new StringTokenizer(br.readLine());
        for (int i = 0; i < n; i++) {
            x[i] = Long.parseLong(st.nextToken());
            x[i + n] = x[i];
        }

        st = new StringTokenizer(br.readLine());
        for (int i = 0; i < n; i++) {
            y[n - i - 1] = Long.parseLong(st.nextToken());
        }

        long[] result = multiply(x, y);

        long maxScore = 0;
        for (int i = n - 1; i < n + n - 1; i++) {
            maxScore = Math.max(maxScore, result[i]);
        }

        System.out.println(maxScore);
    }
}

코드 설명

주요 함수 및 작업

Complex 클래스:

• 복소수 연산을 구현하기 위한 클래스이다.
• FFT 알고리즘에서 사용하는 복소수 곱셈 및 덧셈 연산을 지원한다.

FFT 함수:

• FFT와 역 FFT를 구현한 함수이다.
• invert 플래그를 통해 변환 방향을 제어한다.

배열 곱셈:

• 두 배열 𝑋와 𝑌를 FFT를 활용해 곱한다.
• 결과를 정수로 변환하여 최댓값을 계산한다.

점수 계산:

• 배열 𝑋를 두 배로 확장하여 순환 이동을 처리한다.
• 𝑌는 반전 배열로 설정해 모든 경우를 확인한다.
• 결과 배열에서 최댓값을 구한다.

So...

FFT를 활용하여 𝑂(𝑁log𝑁)의 효율적인 시간 복잡도로 문제를 해결하였다. 순환 이동의 모든 경우를 다루기 위해 𝑋의 확장과 𝑌의 반전을 활용하였고, 복소수 연산과 배열 곱셈의 활용이 핵심이었다. FFT의 구현과 최적화가 중요한 학습 포인트였다.