12월 3일 - 배열의 힘[BOJ/8462]

부분 배열의 힘은 구간 내 모든 숫자에 대해 Ks⋅Ks⋅s의 합으로 정의된다.
여기서 Ks는 해당 숫자가 부분 배열 내에서 등장한 횟수이다.
여러 구간에 대한 쿼리가 주어지며, 각 쿼리에 대해 힘을 효율적으로 계산해야 한다.

코드

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

public class Main {

    static class Query implements Comparable<Query> {
        int left, right, index, block;

        Query(int left, int right, int index, int block) {
            this.left = left;
            this.right = right;
            this.index = index;
            this.block = block;
        }

        @Override
        public int compareTo(Query other) {
            if (this.block != other.block) {
                return Integer.compare(this.block, other.block);
            }
            return Integer.compare(this.right, other.right);
        }
    }

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

        int n = Integer.parseInt(st.nextToken());
        int t = Integer.parseInt(st.nextToken());
        int[] arr = new int[n + 1]; // 1-based indexing
        st = new StringTokenizer(br.readLine());

        for (int i = 1; i <= n; i++) {
            arr[i] = Integer.parseInt(st.nextToken());
        }

        int sqrtN = (int) Math.sqrt(n);
        Query[] queries = new Query[t];

        for (int i = 0; i < t; 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, l / sqrtN);
        }

        Arrays.sort(queries);

        long[] results = new long[t];
        long currentValue = 0;
        int[] count = new int[1000001];
        int currentLeft = 1, currentRight = 0;

        for (Query q : queries) {
            while (currentRight < q.right) {
                currentRight++;
                currentValue += add(arr[currentRight], count);
            }
            while (currentRight > q.right) {
                currentValue -= remove(arr[currentRight], count);
                currentRight--;
            }
            while (currentLeft < q.left) {
                currentValue -= remove(arr[currentLeft], count);
                currentLeft++;
            }
            while (currentLeft > q.left) {
                currentLeft--;
                currentValue += add(arr[currentLeft], count);
            }
            results[q.index] = currentValue;
        }

        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        for (long res : results) {
            bw.write(res + "\n");
        }
        bw.flush();
        bw.close();
    }

    private static long add(int value, int[] count) {
        count[value]++;
        return (long) count[value] * count[value] * value 
             - (long) (count[value] - 1) * (count[value] - 1) * value;
    }

    private static long remove(int value, int[] count) {
        long currentContribution = (long) count[value] * count[value] * value;
        count[value]--;
        return currentContribution 
             - (long) count[value] * count[value] * value;
    }
}

풀이 과정

1. 쿼리 정렬: Mo's Algorithm 사용

   • 각 쿼리를 sqrt-decomposition에 따라 정렬한다.
   • 구간의 왼쪽 (l) 값은 블록 크기로 나누어 정렬하고, 같은 블록 내에서는 오른쪽 (r) 값 기준으로 정렬한다.
   • 이렇게 정렬하면 구간의 이동 횟수를 최소화할 수 있다.

2. 구간 확장 및 축소

   • 현재 [currentLeft, currentRight] 구간을 유지하면서, 각 쿼리의 [l, r]로 확장하거나 축소한다.
   • 숫자가 추가되거나 제거될 때마다 해당 숫자가 구간에 기여하는 힘의 변화량을 계산한다.

3. 숫자 추가 및 제거

   • add(int value, int[] count)
     • 해당 숫자가 구간에 추가될 때, 기여하는 힘을 계산한다.
     • 계산식:
       (현재 등장 횟수)² ⋅ 숫자 − (이전 등장 횟수)² ⋅ 숫자
   • remove(int value, int[] count)
     • 해당 숫자가 구간에서 제거될 때, 기여하는 힘을 감소시킨다.
     • 계산식:
       (이전 등장 횟수)² ⋅ 숫자 − (현재 등장 횟수)² ⋅ 숫자

4. 최적화

   • 구간 이동 시 매번 계산하지 않고, 추가 및 제거 연산만 수행하도록 구현한다.
   • 쿼리의 결과는 정렬된 순서대로 처리되므로, 출력은 원래 순서로 복원한다.

코드 설명

1. Query 클래스

static class Query implements Comparable<Query> {
    int left, right, index, block;

    Query(int left, int right, int index, int block) {
        this.left = left;
        this.right = right;
        this.index = index;
        this.block = block;
    }

    @Override
    public int compareTo(Query other) {
        if (this.block != other.block) {
            return Integer.compare(this.block, other.block);
        }
        return Integer.compare(this.right, other.right);
    }
}

• 쿼리를 정렬하기 위해 사용된다.
• block 기준으로 정렬하고, 같은 블록에서는 right 기준으로 정렬한다.

2. 구간 확장 및 축소

for (Query q : queries) {
    while (currentRight < q.right) {
        currentRight++;
        currentValue += add(arr[currentRight], count);
    }
    while (currentRight > q.right) {
        currentValue -= remove(arr[currentRight], count);
        currentRight--;
    }
    while (currentLeft < q.left) {
        currentValue -= remove(arr[currentLeft], count);
        currentLeft++;
    }
    while (currentLeft > q.left) {
        currentLeft--;
        currentValue += add(arr[currentLeft], count);
    }
    results[q.index] = currentValue;
}

• 구간의 범위를 이동하며, 필요에 따라 숫자를 추가하거나 제거한다.

3. 숫자 추가 및 제거

private static long add(int value, int[] count) {
    count[value]++;
    return (long) count[value] * count[value] * value 
         - (long) (count[value] - 1) * (count[value] - 1) * value;
}

private static long remove(int value, int[] count) {
    long currentContribution = (long) count[value] * count[value] * value;
    count[value]--;
    return currentContribution 
         - (long) count[value] * count[value] * value;
}

• add는 숫자가 추가될 때의 힘 증가량을 계산하며, remove는 제거될 때의 감소량을 계산한다.

So...

이 문제는 구간 쿼리를 효율적으로 처리하는 Mo's Algorithm의 활용 사례이다.
각 숫자의 기여도를 동적으로 업데이트하며, 구간 이동을 최소화하여 효율성을 극대화했다.
가장 큰 고민은 구간 이동과 추가/제거 시의 정확한 힘 계산이었으며, 이를 수식으로 깔끔히 정리하여 해결할 수 있었다.
이 접근법을 통해 O((n + t) √n) 복잡도로 문제를 해결할 수 있었다.