[백준 17435] 합성함수와 쿼리

희소 배열 문제

www.acmicpc.net/problem/17435

희소 배열 기초입니다. 

www.acmicpc.net/problem/11438

응용을 원한다면 위 문제도 풀어보는 것을 권장합니다.

---

[문제 풀이]

 

1부터 m까지의 수에 대해 f(x)가 주어졌을 때, n번 합성한 합성함수를 구하는 문제입니다.

 

걸핏 보면 간단한 문제처럼 보입니다.

m개의 f(x)를 알고 있으니 n번만큼 x = f(x)를 반복하면 가볍게 해결이 가능합니다.

그러나 쿼리의 개수가 200000개, 거슬러 올라가야 하는 n의 제한이 500000이므로 단순히 x = f(x)만 반복한다면 O(nQ)의 시간 복잡도가 나오게 되고, 당연히 시간 초과가 나게 됩니다.

 

즉, 푸는 방법은 간단하지만 시간 복잡도를 향상시켜 풀어야 하는 문제입니다.

어떻게든 시간 복잡도를 줄여야 1초라는 시간안에 통과가 가능할 것 같습니다.

 

우선, 시간 복잡도에 영향을 미치는 조건은 n과 Q입니다. n과 Q를 어떻게든 줄여야 시간 안에 통과가 가능할 것이라는 생각이 듭니다.

하지만 Q는 쿼리의 개수이므로 무슨 방법을 써도 Q를 줄일 수는 없습니다.

그렇다면 남은 조건은 n이고 n을 logn정도로 줄인다면 시간안에 들어올 수 있을 것 같습니다.

 

왜 탐색에서 O(n)의 시간복잡도가 나왔냐 하면 x = f(x)때문입니다.

올라가야 할 n의 최대는 500000인데 느긋하게 1씩 거슬러 올라가고 있었던 것이 문제였습니다.

 

이때 x = f(x) 말고 x = f(f(x))나, x = f(f(f(f(x))))처럼 2의 배수로 한 번에 거슬러 올라갈 수 있다면 시간 복잡도가 O(logn)으로 나올 수 있습니다.

그리고 2의 배수로 한 번에 합성할 수 있게 하는 것이 희소 배열입니다.

 

 

희소 배열은 dp를 사용합니다.

O(logn)이라는 시간 복잡도를 위해 공간 복잡도를 희생하는 방법이라고 생각하면 될 것 같습니다.

2차원 배열을 생성하여, 각 x에 대해 f_(2^i)(x)를 저장하는 방식을 사용합니다.

 

일반 배열과 희소 배열의 차이를 그림으로 설명하면 이해하기보다 나을 겁니다.

 

일반 배열

일반적으로 f(x)를 저장하는 배열입니다. x는 인덱스 i, f(x)는 배열의 값 arr[i]로 선언하고 사용합니다.

1차원 배열이기 때문에 공간 복잡도에서는 이득이지만 1번 합성한 f(x)만 찾는 방식은 너무 많은 시간이 소요됩니다.

1의 7번째 f(x)를 찾기 위해서 1 -> 5 -> 4 -> 2 -> 1 -> 5 -> 4 -> 2라는 비효율적인 과정을 거쳤습니다. [O(n)]

 

 

 

희소 배열

위처럼 2의 배수의 f(x)를 저장하는 배열입니다. x는 인덱스 j, f_(2^i)(x)는 배열의 값 arr[i][j]로 선언하고 사용합니다.

2차원 배열이기 때문에 공간 복잡도를 손해 봤지만 2^i번째 f(x)를 찾을 수 있어 logn의 시간이 소요됩니다.

1의 7번째 f(x)를 찾기 위해서 1 -> 1 -> 4 -> 2라는 과정을 거쳤습니다. [O(logn)]

(f_4(x) = 1, 다시 f_2(1) = 4, 다시 f_1(4) = 2)

 

이런 방식으로 희소 배열을 사용하여 시간 복잡도를 향상할 수 있음을 알았으니 이제 희소 배열을 어떻게 구현하는지, 구현이 끝난 뒤에는 희소 배열을 어떤 식으로 사용해야 위의 예시에서 보인 방법처럼 logn의 시간 복잡도로 빠르게 합성함수를 구할 수 있는지에 대해 설명하겠습니다.

 

 

우선, 희소 배열은 dp를 사용한다고 앞서 말했습니다.

즉 우리는 f(x)에 대한 정보를 알고 있으니 dp[0][x] (f(x))는 채울 수 있습니다.

위의 그림에서도 확인할 수 있듯이 일반 배열의 값과 희소 배열의 0번째 행의 값은 서로 f(x)이므로 동일함을 알 수 있습니다.

그러므로 dp답게 점화식을 사용하여 dp의 i번째 행 dp[i][x]를 채워봅시다. (이때, i행의 의미는 2^i만큼 합성함을 뜻함)

 

우리가 알고 있는 정보가 f(x)에 대한 정보이므로 f_(2^1)(x)를 구하는 방법은 쉽습니다.

f_(2^1)(x)가 dp[1][x]이고 f(f(x))이므로 dp[0][dp[0][x]]로도 생각할 수 있습니다.

즉, dp[1][x] = dp[0][dp[0][x]]이고, 2^i = 2*2^(i-1) 이므로 2의 멱수 형태로 합성함수를 구한다는 것을 이용하여 dp[i][x] = dp[i - 1][dp[i - 1][x]]라는 점화식을 일반화하여 사용할 수 있습니다.

또한 i행이 2^i만큼 합성함을 뜻하고, n <= 500000이므로, i의 크기를 log(500000)으로 둘 필요가 있습니다.

 

 

이제 희소 배열을 어떻게 사용하는지 설명하겠습니다.

위에서 희소 배열에 대한 설명을 간단히 할 때, 7번째 합성함수를 구하는 과정에서 2의 멱수 중 큰 수부터 확인하면서 합성하였습니다.

그리고 이 방법은 마치 10진수에서 2진수를 변환하는 과정과 비슷하다고 생각할 수 있습니다.

2의 멱수만으로 n을 나타내야 하므로 가장 큰 2의 멱수부터 확인해나가면 10진수를 표현할 수 있습니다.

 

설명만으로는 어려우므로 예시를 보여드리겠습니다.

n = 231

n보다 작은 가장 큰 2의 멱수 128 => 128부터 시작

[231] -> [231 - 128 = 103] -> [103 - 64 = 39] -> [39 - 32 = 7] -> [7 - 16 < 0 이므로 7] -> [7 - 8 < 0 이므로 7] -> [7 - 4 = 3] -> [3 - 4 < 0 이므로 3] -> [3 - 2 = 1] -> [1 - 1 = 0]

위와 같은 방식으로 2의 멱수들만으로 logn의 시간 안에 모든 수를 표현할 수 있습니다.

핵심은 가장 큰 2의 멱수부터 확인해야 정확하게 모든 수를 logn의 시간으로 표현할 수 있습니다.

 

위와 같이 희소 배열을 사용하여 O(Qlogn)의 시간 복잡도로 모든 쿼리를 처리할 수 있습니다.

 

---

[소스 코드] / JAVA

 

 

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

public class Main {
    private final static int log = 18;
	
	public static void main(String[] args) throws IOException{
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int N = Integer.parseInt(br.readLine());

        int[][] parent = new int[log + 1][N + 1];
        StringTokenizer st = new StringTokenizer(br.readLine(), " ");
        for (int i = 1; i < N + 1; i++) {
            parent[0][i] = Integer.parseInt(st.nextToken());
        }

        for (int i = 1; i < log + 1; i++) {
            for (int j = 1; j < N + 1; j++) {
                parent[i][j] = parent[i - 1][parent[i - 1][j]];
            }
        }

        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        int Q = Integer.parseInt(br.readLine());
        while (Q-- > 0) {
            st = new StringTokenizer(br.readLine(), " ");
            int n = Integer.parseInt(st.nextToken()), x = Integer.parseInt(st.nextToken());
            for (int i = log; i > -1; i--) {
                int cur = (1 << i);
                if (n >= cur) {
                    x = parent[i][x];
                    n -= cur;
                    if (n == 0) break;
                }
            }

            bw.write(x + "\n");
        }

        bw.flush();
        bw.close();
        br.close();
    }
}