[BOJ/Platinum 2] 백준 8462 배열의 힘(C++)

문제 링크

https://www.acmicpc.net/problem/8462

8462번: 배열의 힘

 

문제

자연수 n개로 이루어진 배열 a1, a2, a3, ..., an이 있다.

 l부터 r까지 부분 배열은 al, al+1, ..., ar이다.

 Ks는 부분 배열 안에 있는 자연수 s의 개수이다.

부분 배열의 힘이란 모든 자연수 s에 대해서, Ks · Ks · s를 합한 값이다.

배열과 부분 배열의 범위가 주어졌을 때, 각 부분 배열의 힘을 구하는 프로그램을 작성하시오.

입력

첫째 줄에 배열의 크기 n과 부분 배열의 개수 t가 주어진다. (1 ≤ n, t ≤ 10^5) 둘째 줄에는 n개의 자연수 ai(1 ≤ ai ≤ 10^6) 가 주어진다.

다음 t개 줄에는 부분 배열의 범위 li와 ri가 주어진다. (1 ≤ li ≤ ri ≤ n)

출력

입력으로 주어지는 각 부분 배열의 힘을 출력한다.

예제 입력 1

8 3
4 3 1 1 1 3 1 2
2 7
1 6
3 8

예제 출력 1

28
25
21

알고리즘 분류

• 평행 분할

풀이

길이가 10만인 수열의 특정 구간 내에서 특정 원소의 개수를 구하는 작업을 2초 안에 10만번이나 수행하기에는 무리가 있어 모스 알고리즘을 통해 해결한다.

우선 M개의 쿼리문을 구간의 시작과 끝, 그리고 쿼리문 자체의 Index 변수를 갖는 구조체로서 벡터에 저장하고, 구간의 시작 Index를 √ N로 나눈 값에 따라 오름차순으로 정렬한다. 만약 그 값이 같은 경우에는 구간의 끝 Index에 따라 오름차순으로 정렬한다.

이제 각 쿼리문마다 해당하는 정답을 구한다. 먼저 첫 번째 쿼리문에 해당하는 구간을 전부 탐색하면서 숫자가 등장하면 먼저 배열의 힘, 즉 해당 숫자의 개수의 제곱 × 해당 숫자를 값에서 뺀 후 등장하는 숫자의 개수를 증가시키고 갱신되는 새로운 배열의 힘을 값에 더한다.

구간을 전부 탐색하고 나온 값이 현재, 즉 첫 번째 쿼리문의 Index번째에 해당하는 정답이 된다.

이제부터는 M-1개의 쿼리문에 해당하는 정답을 전부 구한다. 현재 지정된 구간의 시작과 끝을 현재 쿼리문에 해당하는 구간으로 맞춰가면서 원소를 추가하거나 제외시킨다. 원소를 제외시킬 때에도 마찬가지로 배열의 힘을 값에서 빼고 해당 원소의 개수를 감소시킨 후 갱신되는 새로운 배열의 힘을 값에 더한다. 작업이 모두 끝나면 나오는 값이 현재 쿼리문의 Index번째에 해당하는 정답이 되고, 다음 쿼리문으로 넘어간다.

코드

#include <iostream>
#include <vector>
#include <string>
#include <cmath>
#include <algorithm>
#define MAX 100001
#define LL long long
#define FASTIO cin.tie(0); cout.tie(0); ios::sync_with_stdio(false);

using namespace std;
struct QUERY {
    int Left, Right, Index;
};

int N, T, S, L, R;
LL A[MAX];
vector<QUERY> Query;
int Left, Right;
LL Value, Maximum;
vector<LL> Count;
LL Answer[MAX];

void input() {
    cin >> N >> T;
    for (int i = 0; i < N; i++) {
        cin >> A[i];
        Maximum = max(Maximum, A[i]);
    }
    for (int i = 0; i < T; i++) {
        cin >> L >> R;
        Query.push_back({ L - 1,R - 1,i });
    }
}

bool Comp(QUERY A, QUERY B) {
    if (A.Left / S != B.Left / S) {
        return (A.Left / S < B.Left / S);
    }
    return (A.Right < B.Right);
}

void settings() {
    S = sqrt(N);
    sort(Query.begin(), Query.end(), Comp);
    Count.resize(Maximum + 1);
    Left = Query[0].Left;
    Right = Query[0].Right;
    for (int i = Left; i <= Right; i++) {
        LL Now = A[i];
        Value -= (Count[Now] * Count[Now] * Now);
        Count[Now]++;
        Value += (Count[Now] * Count[Now] * Now);
    }
    Answer[Query[0].Index] = Value;
    for (int i = 1; i < T; i++) {
        while (Query[i].Left < Left) {
            Left--;
            LL Now = A[Left];
            Value -= (Count[Now] * Count[Now] * Now);
            Count[Now]++;
            Value += (Count[Now] * Count[Now] * Now);
        };

        while (Query[i].Left > Left) {
            LL Now = A[Left];
            Value -= (Count[Now] * Count[Now] * Now);
            Count[Now]--;
            Value += (Count[Now] * Count[Now] * Now);
            Left++;
        };

        while (Query[i].Right > Right) {
            Right++;
            LL Now = A[Right];
            Value -= (Count[Now] * Count[Now] * Now);
            Count[Now]++;
            Value += (Count[Now] * Count[Now] * Now);
        };

        while (Query[i].Right < Right) {
            LL Now = A[Right];
            Value -= (Count[Now] * Count[Now] * Now);
            Count[Now]--;
            Value += (Count[Now] * Count[Now] * Now);
            Right--;
        };

        Answer[Query[i].Index] = Value;
    }
}

void find_Answer() {
    for (int i = 0; i < T; i++) {
        cout << Answer[i] << "\n";
    }
}

int main() {
    FASTIO
    input();
    settings();
    find_Answer();

    return 0;
}