문제

아이디어

문제에서 주어진 데이터를 일차원 배열 arr에 저장한다고 가정하자. $f_n(x)$을 구하는 메서드 query(n, x)는 아래와 같이 구현할 수 있다.

static int query(int n, int x) {
    for (int i = 0; i < n; i++) {
        x = arr[x];
    }
    return x;
}

간단한 구현이지만 시간복잡도는 O(Q * N)에 달한다. 주어진 시간 제한 내에는 풀이할 수 없는 방법이다.

그렇다면 주어진 데이터를 이차원 배열 arr에 저장한다고 가정하자. 이 때, $f_n(x) = arr[x][n]$으로 저장된다.

하지만 arr의 공간 복잡도는 O(N * M)에 달한다. 정수형 배열일 때 무려 400GB 이상의 크기를 차지한다. 물론 불가능한 크기다.

2차원 배열을 사용하는 풀이의 문제점은 과도한 메모리를 할당할 뿐만 아니라 배열의 전부를 사용하지 않는다는 점이다. 최대 100억개의 원소가 할당되지만 실제 사용되는 원소는 고작 20만개를 넘지 않는다.

여기서 희소 배열의 필요성을 찾을 수 있다.

희소 배열

대부분의 요소가 비어있거나 0인 배열을 효율적으로 저장하는 자료구조

2차원 배열 arr를 $f_{2^n}(x) = arr[x][n]$으로 재정의하자. 그렇다면 메서드 query(n, x)는 아래와 같이 사용할 수 있다.

static int query(int n, int x) {
    for (int i = (int) Math.sqrt(n); i >= 0; i--) {
        if ((n & (1 << i)) == 0)
            continue;
        x = sparseTable[x][i];
    }
    return x;
}

소스코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.StringTokenizer;

public class Main {
    static int M, Q;

    static int[][] sparseTable;

    static void setTable() {
        for (int j = 1; j < 19; j++) {
            for (int i = 1; i <= M; i++) {
                sparseTable[i][j] = sparseTable[sparseTable[i][j - 1]][j - 1];
            }
        }
    }

    static int query(int n, int x) {
        for (int i = 18; i >= 0; i--) {
            if ((n & (1 << i)) == 0)
                continue;
            x = sparseTable[x][i];
        }
        return x;
    }

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringBuilder ans = new StringBuilder();
        StringTokenizer st;

        M = Integer.parseInt(br.readLine());

        sparseTable = new int[M + 1][19];
        st = new StringTokenizer(br.readLine());
        for (int i = 1; i <= M; i++) {
            sparseTable[i][0] = Integer.parseInt(st.nextToken());
        }

        setTable();

        Q = Integer.parseInt(br.readLine());
        for (int i = 0; i < Q; i++) {
            st = new StringTokenizer(br.readLine());

            int n = Integer.parseInt(st.nextToken());
            int x = Integer.parseInt(st.nextToken());

            ans.append(query(n, x)).append('\n');
        }

        System.out.print(ans);
    }
}

채점결과