11월 25일 - 수열과 쿼리 3 [BOJ/13544]

이 문제는 수열에서 특정 구간 내에서 k보다 큰 원소의 개수를 빠르게 구하는 것이다. 쿼리의 조건으로 XOR 연산을 통해 입력이 변동되기 때문에 효율적인 데이터 구조 설계가 필요했다. 머지 소트 트리(Merge Sort Tree)를 활용해 구간 데이터를 정렬된 상태로 유지하며, 쿼리를 O(log^2 N)에 처리할 수 있었다.

코드

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

class Main {

    static int N, M, lastAnswer = 0;
    static ArrayList<Integer>[] tree; // 세그먼트 트리
    static int[] arr;

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

        // 입력
        N = Integer.parseInt(br.readLine());
        arr = new int[N];
        st = new StringTokenizer(br.readLine());
        for (int i = 0; i < N; i++) {
            arr[i] = Integer.parseInt(st.nextToken());
        }

        // 세그먼트 트리 초기화
        tree = new ArrayList[4 * N];
        for (int i = 0; i < 4 * N; i++) {
            tree[i] = new ArrayList<>();
        }
        init(1, 0, N - 1);

        // 쿼리 처리
        M = Integer.parseInt(br.readLine());
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));

        for (int i = 0; i < M; i++) {
            st = new StringTokenizer(br.readLine());
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());
            int c = Integer.parseInt(st.nextToken());

            // XOR로 i, j, k 계산
            int i1 = (a ^ lastAnswer);
            int j1 = (b ^ lastAnswer);
            int k1 = (c ^ lastAnswer);

            // i와 j를 1-based index에서 0-based index로 변환
            i1--;
            j1--;

            // 쿼리 결과 계산
            lastAnswer = query(1, 0, N - 1, i1, j1, k1);
            bw.write(lastAnswer + "\n");
        }

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

    static void init(int node, int start, int end) {
        if (start == end) {
            tree[node].add(arr[start]);
        } else {
            int mid = (start + end) / 2;
            init(node * 2, start, mid);
            init(node * 2 + 1, mid + 1, end);

            // 두 자식 노드의 데이터를 병합
            merge(tree[node * 2], tree[node * 2 + 1], tree[node]);
        }
    }

    static void merge(ArrayList<Integer> left, ArrayList<Integer> right, ArrayList<Integer> result) {
        int i = 0, j = 0;
        while (i < left.size() && j < right.size()) {
            if (left.get(i) <= right.get(j)) {
                result.add(left.get(i++));
            } else {
                result.add(right.get(j++));
            }
        }
        while (i < left.size()) result.add(left.get(i++));
        while (j < right.size()) result.add(right.get(j++));
    }

    static int query(int node, int start, int end, int L, int R, int val) {
        if (end < L || R < start) {
            return 0; // 범위를 벗어난 경우
        }
        if (L <= start && end <= R) {
            // 현재 구간이 [L, R]에 완전히 포함되는 경우
            return tree[node].size() - upperBound(tree[node], val);
        }
        int mid = (start + end) / 2;
        int leftResult = query(node * 2, start, mid, L, R, val);
        int rightResult = query(node * 2 + 1, mid + 1, end, L, R, val);
        return leftResult + rightResult;
    }

    static int upperBound(ArrayList<Integer> list, int value) {
        int low = 0, high = list.size();
        while (low < high) {
            int mid = (low + high) / 2;
            if (list.get(mid) > value) {
                high = mid;
            } else {
                low = mid + 1;
            }
        }
        return low;
    }
}

코드 설명

1. 세그먼트 트리 초기화

static void init(int node, int start, int end) {
    if (start == end) {
        tree[node].add(arr[start]);
    } else {
        int mid = (start + end) / 2;
        init(node * 2, start, mid);
        init(node * 2 + 1, mid + 1, end);

        // 두 자식 노드의 데이터를 병합
        merge(tree[node * 2], tree[node * 2 + 1], tree[node]);
    }
}

• 리프 노드에는 각 배열의 값을 저장한다.
• 내부 노드는 두 자식 노드를 정렬된 상태로 병합해 구간 내 데이터가 항상 정렬 상태를 유지하도록 한다.

2. 병합 함수

static void merge(ArrayList<Integer> left, ArrayList<Integer> right, ArrayList<Integer> result) {
    int i = 0, j = 0;
    while (i < left.size() && j < right.size()) {
        if (left.get(i) <= right.get(j)) {
            result.add(left.get(i++));
        } else {
            result.add(right.get(j++));
        }
    }
    while (i < left.size()) result.add(left.get(i++));
    while (j < right.size()) result.add(right.get(j++));
}

• 병합 정렬(Merge Sort)의 병합 단계와 동일하다.
• 각 구간 내 데이터를 정렬된 상태로 유지해 이분 탐색이 가능하도록 한다.

3. 쿼리 처리

static int query(int node, int start, int end, int L, int R, int val) {
    if (end < L || R < start) {
        return 0; // 범위를 벗어난 경우
    }
    if (L <= start && end <= R) {
        // 현재 구간이 [L, R]에 완전히 포함되는 경우
        return tree[node].size() - upperBound(tree[node], val);
    }
    int mid = (start + end) / 2;
    int leftResult = query(node * 2, start, mid, L, R, val);
    int rightResult = query(node * 2 + 1, mid + 1, end, L, R, val);
    return leftResult + rightResult;
}

• 특정 구간 [L, R]에 대해 k보다 큰 원소의 개수를 구한다.
• upperBound를 통해 val보다 큰 값의 첫 번째 위치를 찾고, 이를 기반으로 결과를 계산한다.

4. upperBound 구현

static int upperBound(ArrayList<Integer> list, int value) {
    int low = 0, high = list.size();
    while (low < high) {
        int mid = (low + high) / 2;
        if (list.get(mid) > value) {
            high = mid;
        } else {
            low = mid + 1;
        }
    }
    return low;
}

• 이분 탐색을 통해 val보다 큰 값이 처음 나타나는 위치를 찾는다.

So...

머지 소트 트리를 활용해 효율적으로 구간 내 k보다 큰 값의 개수를 구하는 방법을 다루었다. 데이터가 정렬된 상태로 유지되므로, 구간 내 쿼리를 처리할 때 이분 탐색으로 빠르게 결과를 계산할 수 있었다. XOR 연산으로 동적으로 입력이 변하는 조건을 처리하기 위해 추가적인 계산을 수행했으며, 쿼리의 복잡도는 O(log^2 N)으로 유지할 수 있었다. 문제를 풀면서 구간 데이터를 정렬 상태로 유지하는 중요성을 다시 한번 느낄 수 있었다.