11월 19일 - 수열과 쿼리 1 [BOJ/13537]

이번 문제는 구간 내에서 특정 값보다 큰 원소의 개수를 효율적으로 계산하는 것이다. 배열의 길이와 쿼리의 수가 최대 100,000까지 가능하기 때문에, 효율적인 자료구조와 알고리즘이 필요했다. 이를 위해 머지 소트 트리(Merge Sort Tree)를 사용해 해결했다.

---

문제 접근 방법

1. 머지 소트 트리의 개념:

• 머지 소트 트리는 세그먼트 트리와 병합 정렬의 아이디어를 결합한 자료구조다.
• 각 노드에는 해당 구간의 원소들이 정렬된 상태로 저장된다.
• 특정 구간에 대해 원소들이 정렬된 상태를 유지하기 때문에, 이진 탐색을 통해 구간 내 조건을 만족하는 원소의 개수를 빠르게 구할 수 있다.

2. 쿼리 처리 로직:

• 쿼리 (i, j, k)에 대해 i부터 j까지의 구간에서 k보다 큰 원소의 개수를 구해야 한다.
• 머지 소트 트리를 사용해 해당 구간의 노드들을 합치고, 이 합쳐진 리스트에서 k보다 큰 원소의 개수를 upper_bound를 통해 계산한다.

3. 시간 복잡도 분석:

• 트리 빌드: O(N log N)
  • 트리를 빌드할 때, 각 노드에 정렬된 리스트를 저장하므로 O(N log N)이다.
• 쿼리 처리: O(log N * log N)
  • 세그먼트 트리 탐색(O(log N))과 정렬된 리스트에서 이진 탐색(O(log N))이 결합된 구조다.
• 따라서, 전체 쿼리 처리 시간은 O((N + M) log^2 N)이다.

---

코드

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

public class Main {
    static int N, M, k;
    static int[] A;
    static List<Integer>[] segmentTree;

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

        // 입력 처리
        N = Integer.parseInt(br.readLine());
        A = new int[N + 1];
        st = new StringTokenizer(br.readLine());
        for (int i = 1; i <= N; i++) {
            A[i] = Integer.parseInt(st.nextToken());
        }
        M = Integer.parseInt(br.readLine());

        // 머지 소트 트리 초기화
        segmentTree = new ArrayList[4 * N];
        for (int i = 0; i < 4 * N; i++) {
            segmentTree[i] = new ArrayList<>();
        }
        build(1, N, 1);

        // 쿼리 처리
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < M; i++) {
            st = new StringTokenizer(br.readLine());
            int l = Integer.parseInt(st.nextToken());
            int r = Integer.parseInt(st.nextToken());
            k = Integer.parseInt(st.nextToken());

            int result = query(1, N, 1, l, r);
            sb.append(result).append("\n");
        }

        // 결과 출력
        System.out.print(sb.toString());
    }

    // 머지 소트 트리 빌드
    static void build(int start, int end, int node) {
        if (start == end) {
            segmentTree[node].add(A[start]);
        } else {
            int mid = (start + end) / 2;
            build(start, mid, node * 2);
            build(mid + 1, end, node * 2 + 1);

            segmentTree[node] = merge(segmentTree[node * 2], segmentTree[node * 2 + 1]);
        }
    }

    // 두 리스트 병합
    static List<Integer> merge(List<Integer> left, List<Integer> right) {
        List<Integer> merged = new ArrayList<>();
        int i = 0, j = 0;

        while (i < left.size() && j < right.size()) {
            if (left.get(i) <= right.get(j)) {
                merged.add(left.get(i++));
            } else {
                merged.add(right.get(j++));
            }
        }

        while (i < left.size()) merged.add(left.get(i++));
        while (j < right.size()) merged.add(right.get(j++));

        return merged;
    }

    // 쿼리 처리
    static int query(int start, int end, int node, int left, int right) {
        if (right < start || end < left) {
            return 0;
        }

        if (left <= start && end <= right) {
            return segmentTree[node].size() - upperBound(segmentTree[node], k);
        }

        int mid = (start + end) / 2;
        return query(start, mid, node * 2, left, right) +
               query(mid + 1, end, node * 2 + 1, left, right);
    }

    // upper_bound 구현
    static int upperBound(List<Integer> list, int value) {
        int low = 0, high = list.size();

        while (low < high) {
            int mid = (low + high) / 2;
            if (list.get(mid) <= value) {
                low = mid + 1;
            } else {
                high = mid;
            }
        }

        return low;
    }
}

코드 설명

1. 입력 처리와 초기화

N = Integer.parseInt(br.readLine());
A = new int[N + 1];
st = new StringTokenizer(br.readLine());
for (int i = 1; i <= N; i++) {
    A[i] = Integer.parseInt(st.nextToken());
}
M = Integer.parseInt(br.readLine());
segmentTree = new ArrayList[4 * N];
for (int i = 0; i < 4 * N; i++) {
    segmentTree[i] = new ArrayList<>();
}

• 수열과 쿼리의 입력을 처리한다.
• 세그먼트 트리의 각 노드를 리스트 형태로 초기화한다.

2. 머지 소트 트리 빌드

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

        segmentTree[node] = merge(segmentTree[node * 2], segmentTree[node * 2 + 1]);
    }
}

• 리프 노드에는 해당 구간의 값을 저장한다.
• 내부 노드는 자식 노드의 리스트를 병합하여 정렬된 상태로 저장한다.

3. 두 리스트 병합

static List<Integer> merge(List<Integer> left, List<Integer> right) {
    List<Integer> merged = new ArrayList<>();
    int i = 0, j = 0;

    while (i < left.size() && j < right.size()) {
        if (left.get(i) <= right.get(j)) {
            merged.add(left.get(i++));
        } else {
            merged.add(right.get(j++));
        }
    }

    while (i < left.size()) merged.add(left.get(i++));
    while (j < right.size()) merged.add(right.get(j++));

    return merged;
}

• 두 리스트를 병합 정렬 방식으로 병합한다.

4. 쿼리 처리

static int query(int start, int end, int node, int left, int right) {
    if (right < start || end < left) {
        return 0;
    }

    if (left <= start && end <= right) {
        return segmentTree[node].size() - upperBound(segmentTree[node], k);
    }

    int mid = (start + end) / 2;
    return query(start, mid, node * 2, left, right) +
           query(mid + 1, end, node * 2 + 1, left, right);
}

• 구간이 쿼리에 완전히 포함되면 upper_bound를 사용해 k보다 큰 원소의 개수를 계산한다.
• 구간이 겹치는 경우 재귀적으로 탐색한다.

5. 이진 탐색 (upper_bound)

static int upperBound(List<Integer> list, int value) {
    int low = 0, high = list.size();

    while (low < high) {
        int mid = (low + high) / 2;
        if (list.get(mid) <= value) {
            low = mid + 1;
        } else {
            high = mid;
        }
    }

    return low;
}

• 정렬된 리스트에서 value보다 큰 첫 번째 위치를 반환한다.

So...

머지 소트 트리를 사용해 구간 내 조건을 만족하는 원소 개수를 효율적으로 계산할 수 있었고, 최적화된 방식으로 모든 쿼리를 처리할 수 있었다.