BOJ 10800 : 컬러볼 [C++]

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

 

공의 개수가 최대 20만 개, 색상 수도 최대 20만 개, 크기는 최대 2천 가지라 어떻게 해결해야 할지 고민을 좀 많이 했다.

크기 순으로 정렬을 하더라도, 같은 색상의 공들은 빼야 하고, 같은 크기를 가지는 공들도 제외해야 한다.

 

이를 해결하기 위해 같은 색상 공들의 크기 합, 같은 크기를 가지는 공들의 크기 합을 따로 보관해야겠다고 생각했다.

같은 색상 공들의 크기를 따로 보관하면, 전체 공 크기의 합에서 같은 색상 공들의 크기 합을 빼 해당 공의 점수를 구할 수 있게 된다.

같은 크기를 가지는 공은 순회하면서 앞의 공과 크기가 달라질 때마다 비워주면서 전체 공의 크기 합, 같은 색상 공들의 크기합을 누적하도록 했다.

 

 

로직 순서

1. 색상 별 크기의 합을 유지하는 배열 sumPerColor를 만들어 같은 색상의 공들은 뺄 수 있도록 했다.

2. 크기 순으로 공을 정렬

3. 크기 순으로 정렬된 배열을 순회

3 - 1. 앞의 공과 크기가 다르다면, sameSize 벡터를 순회하며 지금까지 살펴본 공 크기의 총합(sum)에 공의 크기를 더하고, 공의 색상별로 sumPerColor를 더해주고 sameSize 벡터를 비운다.

3 - 2. answer[index]에 sum - sumPerColor[color]를 넣고, 해당 공의 정보를 sameSize 벡터에 보관.

4. 정답 출력

 

코드

#include <iostream>
#include <vector>
#include <algorithm>

struct ball {
    int index;
    int color;
    int size;
    ball() {}
    ball(int index, int color, int size) : index(index), color(color), size(size) {}
};

int N;
ball balls[200001];
int answer[200001];
int sum = 0;
int sumPerColor[200001];

int main() {
    std::cin.tie(NULL);
    std::ios::sync_with_stdio(false);

    std::cin >> N;

    int color, size;
    int maxColor = 0;
    for (int i = 0; i < N; i++) {
        std::cin >> color >> size;
        balls[i + 1] = ball(i + 1, color, size);
        maxColor = std::max(color, maxColor);
    }

    std::sort(balls + 1, balls + N + 1, [&](ball a, ball b) {
        if (a.size == b.size) {
            return a.color > b.color;
        } else {
            return a.size < b.size;
        }
    });

    balls[0] = ball(0, 0, 0);
    std::vector<std::pair<int, int>> sameSize;
    for (int i = 1; i < N + 1; i++) {
        if (balls[i].size != balls[i - 1].size) {
            for (auto p : sameSize) {
                sum += p.second;
                sumPerColor[p.first] += p.second;
            }
            sameSize.clear();
        }
        sameSize.push_back({balls[i].color, balls[i].size});
        answer[balls[i].index] = sum - sumPerColor[balls[i].color];
    }

    for (int i = 1; i < N + 1; i++) {
        std::cout << answer[i] << "\n";
    }

    return 0;
}