TIL: Mo's Algorithm을 활용한 최댓값 쿼리 처리

문제 설명

길이가 N인 수열이 주어진다. 수열의 원소는 1 ≤ Ai ≤ K를 만족하며,
다음과 같은 쿼리를 M번 수행해야 한다.

• 쿼리 (l, r):
  • 수열 A[l] ~ A[r]에서 같은 값을 가지는 x, y의 인덱스 간 최대 거리를 구한다.
  • 즉, max{|x − y| : l ≤ x, y ≤ r and Ax = Ay}를 출력한다.

해결 방법

1. TreeSet을 활용한 거리 계산

각 숫자에 대해 해당 값이 등장한 인덱스를 TreeSet 에 저장하면,
최댓값을 O(1)에 얻을 수 있다.

• positions[A[i]] = {i1, i2, ..., im} 으로 인덱스를 저장.
• 같은 값이 여러 개 등장할 경우,
  • positions[A[i]].last() - positions[A[i]].first() 로 최대 거리 계산.

2. Mo’s Algorithm 활용

Mo’s Algorithm은 오프라인 쿼리 처리에 적합하며, O(√N * Q) 에 수행 가능하다.

1. 쿼리 정렬

   • SQRT 크기로 블록을 나누고, (L, R) 기준으로 정렬하여 추가 연산을 줄인다.
   • 같은 블록 내에서는 R 값을 기준으로 정렬한다.

2. 투 포인터를 이용해 구간을 조정

   • 현재 L과 R을 유지하며 새로운 구간을 추가/제거하여 빠르게 결과를 갱신한다.

코드

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

public class Main {
    static final int MAX_N = 100000;
    static final int SQRT = 316;
    static int n, k, q;
    static int[] arr = new int[MAX_N + 1];
    static Query[] queries;
    static TreeSet<Integer>[] positions;
    static int[] dcnt = new int[MAX_N + 1], sqcnt = new int[(MAX_N / SQRT) + 10], ans;

    static class Query implements Comparable<Query> {
        int l, r, index;
        Query(int l, int r, int index) {
            this.l = l;
            this.r = r;
            this.index = index;
        }
        @Override
        public int compareTo(Query other) {
            if (l / SQRT != other.l / SQRT) return Integer.compare(l, other.l);
            return Integer.compare(r, other.r);
        }
    }

    static void remove(int idx) {
        int num = arr[idx];
        if (positions[num].size() >= 2) {
            int diff = positions[num].last() - positions[num].first();
            dcnt[diff]--;
            sqcnt[diff / SQRT]--;
        }
        positions[num].remove(idx);
        if (positions[num].size() >= 2) {
            int diff = positions[num].last() - positions[num].first();
            dcnt[diff]++;
            sqcnt[diff / SQRT]++;
        }
    }

    static void add(int idx) {
        int num = arr[idx];
        if (positions[num].size() >= 2) {
            int diff = positions[num].last() - positions[num].first();
            dcnt[diff]--;
            sqcnt[diff / SQRT]--;
        }
        positions[num].add(idx);
        if (positions[num].size() >= 2) {
            int diff = positions[num].last() - positions[num].first();
            dcnt[diff]++;
            sqcnt[diff / SQRT]++;
        }
    }

    static int getMaxDistance() {
        for (int i = sqcnt.length - 1; i >= 0; i--) {
            if (sqcnt[i] > 0) {
                for (int j = SQRT - 1; j >= 0; j--) {
                    int diff = i * SQRT + j;
                    if (diff < dcnt.length && dcnt[diff] > 0) return diff;
                }
            }
        }
        return 0;
    }

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        n = Integer.parseInt(st.nextToken());
        k = Integer.parseInt(st.nextToken());
        st = new StringTokenizer(br.readLine());
        for (int i = 1; i <= n; i++) {
            arr[i] = Integer.parseInt(st.nextToken());
        }
        q = Integer.parseInt(br.readLine());
        queries = new Query[q];
        ans = new int[q];
        positions = new TreeSet[k + 1];
        for (int i = 1; i <= k; i++) positions[i] = new TreeSet<>();

        for (int i = 0; i < q; i++) {
            st = new StringTokenizer(br.readLine());
            int l = Integer.parseInt(st.nextToken());
            int r = Integer.parseInt(st.nextToken());
            queries[i] = new Query(l, r, i);
        }

        Arrays.sort(queries);
        int l = queries[0].l, r = queries[0].r;
        for (int i = l; i <= r; i++) add(i);
        ans[queries[0].index] = getMaxDistance();

        for (int i = 1; i < q; i++) {
            while (queries[i].l < l) add(--l);
            while (r < queries[i].r) add(++r);
            while (l < queries[i].l) remove(l++);
            while (queries[i].r < r) remove(r--);
            ans[queries[i].index] = getMaxDistance();
        }

        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        for (int i = 0; i < q; i++) bw.write(ans[i] + "\n");
        bw.flush();
    }
}

핵심 요약

1. TreeSet을 이용한 최대 거리 계산

   • 각 값의 등장 위치를 TreeSet에 저장하여 ( O(1) )에 최대 거리 조회 가능.

2. Mo’s Algorithm 활용

   • 블록 정렬 후, 투 포인터를 이용해 ( O(\sqrt{N} * Q) ) 로 빠른 쿼리 처리.

3. 오프라인 쿼리 처리

   • 먼저 쿼리를 정렬한 뒤, ( L, R )을 이동하며 새로운 값만 추가/삭제.

So...

이 문제를 통해 Mo’s Algorithm을 활용한 최댓값 쿼리 처리를 배울 수 있었다.
특히, TreeSet을 사용한 최대 거리 계산과 오프라인 쿼리 최적화 기법이 핵심이었다.