TIL: 부분합과 Mo's Algorithm을 활용한 부분 수열 최적화 문제 해결

문제 설명

길이가 N인 1과 -1로 이루어진 수열 A가 주어진다.
여기서 주어지는 M개의 쿼리에 대해 다음과 같은 연산을 수행해야 한다.

• i j : A[i] ~ A[j]로 이루어진 부분 수열 중에서 합이 0이면서 가장 긴 연속한 부분 수열의 길이를 출력한다.

제한 조건

• 1 ≤ N, M ≤ 100,000 으로 매우 크기 때문에 O(N^2) 이상의 알고리즘을 사용할 수 없다.

해결 방법

1. 누적 합(Prefix Sum) 사용

주어진 배열에서 특정 부분 수열의 합을 빠르게 구하기 위해 누적 합(Prefix Sum) 을 활용한다.

• prefix[i] = A[1] + A[2] + ... + A[i]
• A[l] ~ A[r]의 합은 prefix[r] - prefix[l-1] 로 O(1)에 계산 가능하다.

누적 합을 이용한 핵심 아이디어

1. 부분합을 이용해 A[i] ~ A[j]의 합이 0이 되는 구간을 찾는다.
2. prefix[j] = prefix[i] 가 같은 두 i, j에 대해 가장 긴 j - i 값을 찾는다.

2. Mo’s Algorithm 활용

Mo’s Algorithm은 오프라인 쿼리 처리에 적합한 알고리즘으로, O(√N * Q) 에 연산이 가능하다.

Mo's Algorithm의 원리

1. 쿼리 정렬

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

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

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

코드

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

public class Main {
    static final int SQRT = 400;
    static final int MAX_N = 100000;
    static int n, q;
    static int[] arr = new int[MAX_N + 1];
    static List<Query> queries = new ArrayList<>();
    static List<Integer>[] pos = new ArrayList[MAX_N * 2 + 1];
    static int[] count = new int[MAX_N + 1], bucket = new int[(MAX_N / SQRT) + 10], ans;

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

    static void add(int x, int dir) {
        int current = 0;
        List<Integer> list = pos[arr[x]];
        if (!list.isEmpty()) {
            current = list.get(list.size() - 1) - list.get(0);
            count[current]--;
            bucket[current / SQRT]--;
        }
        if (dir == 0) list.add(0, x);
        else list.add(x);
        current = list.get(list.size() - 1) - list.get(0);
        count[current]++;
        bucket[current / SQRT]++;
    }

    static void remove(int x, int dir) {
        List<Integer> list = pos[arr[x]];
        int current = list.get(list.size() - 1) - list.get(0);
        count[current]--;
        bucket[current / SQRT]--;
        if (dir == 0) list.remove(0);
        else list.remove(list.size() - 1);
        if (!list.isEmpty()) {
            current = list.get(list.size() - 1) - list.get(0);
            count[current]++;
            bucket[current / SQRT]++;
        }
    }

    static int query() {
        for (int i = bucket.length - 1; i >= 0; i--) {
            if (bucket[i] == 0) continue;
            for (int j = SQRT - 1; j >= 0; j--) {
                int index = i * SQRT + j;
                if (index < count.length && count[index] > 0) return index;
            }
        }
        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());
        st = new StringTokenizer(br.readLine());
        for (int i = 1; i <= n; i++) {
            arr[i] = Integer.parseInt(st.nextToken());
        }
        for (int i = 1; i <= n; i++) {
            arr[i] += arr[i - 1];
        }
        for (int i = 0; i <= n; i++) {
            arr[i] += MAX_N;
        }
        for (int i = 0; i < pos.length; i++) pos[i] = new ArrayList<>();
        st = new StringTokenizer(br.readLine());
        q = Integer.parseInt(st.nextToken());
        ans = new int[q];
        for (int i = 0; i < q; i++) {
            st = new StringTokenizer(br.readLine());
            int s = Integer.parseInt(st.nextToken()) - 1;
            int e = Integer.parseInt(st.nextToken());
            queries.add(new Query(s, e, i));
        }
        queries.sort(null);
        int s = queries.get(0).s, e = queries.get(0).e, x = queries.get(0).index;
        for (int i = s; i <= e; i++) add(i, 1);
        ans[x] = query();
        for (int i = 1; i < q; i++) {
            x = queries.get(i).index;
            while (queries.get(i).s < s) add(--s, 0);
            while (e < queries.get(i).e) add(++e, 1);
            while (s < queries.get(i).s) remove(s++, 0);
            while (queries.get(i).e < e) remove(e--, 1);
            ans[x] = query();
        }
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        for (int i = 0; i < q; i++) bw.write(ans[i] + "\n");
        bw.flush();
    }
}

핵심 요약

1. 누적 합 (Prefix Sum)

   • ( A[L] \sim A[R] ) 의 합을 빠르게 계산한다.

2. 해시 테이블 활용

   • 해시 테이블을 사용하여 동일한 prefix[j] 값을 가진 가장 먼 (i, j)를 찾는다.

3. Mo's Algorithm 적용

   • ( O(\sqrt{N} * Q) ) 의 시간 복잡도로 쿼리를 처리한다.

4. 오프라인 쿼리 처리

   • 빠르게 정렬 후, 투 포인터를 이용해 효율적으로 ( L, R )을 이동시킨다.

So...

이 문제를 통해 누적 합과 Mo’s Algorithm 을 결합하여 빠른 쿼리 처리 방법을 익혔다.
특히, Mo's Algorithm의 블록 정렬과 투 포인터 활용법을 실전에 적용할 수 있었다.