백준 컬러볼 [c++]

10800 컬러볼 골드 2

문제 링크

문제

지훈이가 최근에 즐기는 컴퓨터 게임이 있다. 이 게임은 여러 플레이어가 참여하며, 각 플레이어는 특정한 색과 크기를 가진 자기 공 하나를 조종하여 게임에 참여한다. 각 플레이어의 목표는 자기 공보다 크기가 작고 색이 다른 공을 사로잡아 그 공의 크기만큼의 점수를 얻는 것이다. 그리고 다른 공을 사로잡은 이후에도 본인의 공의 색과 크기는 변하지 않는다. 다음 예제는 네 개의 공이 있다. 편의상 색은 숫자로 표현한다.

공 번호 색 크기

Index Color Size
1 1 10
2 3 15
3 1 3
4 4 8

이 경우, 2번 공은 다른 모든 공을 사로잡을 수 있다. 반면, 1번 공은 크기가 더 큰 2번 공과 색이 같은 3번 공은 잡을 수 없으며, 단지 4번 공만 잡을 수 있다.

공들의 색과 크기가 주어졌을 때, 각 플레이어가 사로잡을 수 있는 모든 공들의 크기의 합을 출력하는 프로그램을 작성하시오.

입력

첫 줄에는 공의 개수를 나타내는 자연수 N이 주어진다(1 ≤ N ≤ 200,000). 다음 N개의 줄 중 i번째 줄에는 i번째 공의 색을 나타내는 자연수 Ci와 그 크기를 나타내는 자연수 Si가 주어진다(1 ≤ Ci ≤ N, 1 ≤ Si ≤ 2,000). 서로 같은 크기 혹은 같은 색의 공들이 있을 수 있다.

출력

N개의 줄을 출력한다. N개의 줄 중 i번째 줄에는 i번째 공을 가진 플레이어가 잡을 수 있는 모든 공들의 크기 합을 출력한다.

풀이

처음 접근한 방법은, 크기와 색을 백터에 저장한 후 크기 순으로 정렬하고나서 하나씩 돌며 자신보다 크기가 작거나 색이 다른 값들을 더해가는 방법으로 구현했다.

당연히 삼중for문을 사용해 시간초과가 났고 이중 for문으로 줄여봐도 마찬가지였다.

문제를 다시 읽고 생각해보니, 크기 순으로 정렬하고나면 자신보다 크기가 작은 값들의 합을 계산하는데 이는 다음 인덱스의 값도 똑같은 과정을 하므로 부분합을 이용할 수 있었다.

최종 sum = (지금 인덱스까지 부분합) - (자신과 색이 같은 공들의 부분합) - (나와 크키가 같은 공들의 부분합) + 내 크기 로 결정된다.

예외 처리해야할 부분이 있는데, 만약 크기순으로 정렬했을때, 크기와 색이 모두 같은 경우, 이전 인덱스 값과 동일하다.

코드

#include <iostream>
#include <algorithm>
#include <vector>
#include <string.h>
#include <map>
using namespace std;

typedef struct info {
    int idx;
    int color;
    int size;
} info;

vector<info> input;
vector<long long> ans(200001,0);
int n;

bool comp(info a, info b)
{
    if(a.size == b.size) return a.color < b.color;
    return a.size < b.size;
}

int main(){
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);
    cin >> n;

    for(int i = 0; i < n; i++)
    {
        int color;
        int size;
        cin >> color >> size;
        info inf;
        inf.idx = i;
        inf.color = color;
        inf.size = size;
        input.push_back(inf);
        // box[color - 1].push_back(size);
    }

    sort(input.begin(), input.end(), comp);

    long long sum = 0;
    // long long sum_color[200001];
    // long long sum_size[200001];
    // memset(sum_color, 0, 200001);
    // memset(sum_size, 0, 200001);
    vector<int> sum_color (200001, 0);
    vector<int> sum_size (2001, 0);
    for(int i = 0; i < n; i++)
    {
        int index = input[i].idx;
        int color = input[i].color;
        int size = input[i].size;
        sum += size;
        sum_color[color] += size;
        sum_size[size] += size;

        if(i > 0 && size == input[i-1].size && color == input[i-1].color){
            ans[index] = ans[input[i-1].idx];
        } else{
            ans[index] = sum - sum_color[color] - sum_size[size] + size;
        }
    }

    for(int i = 0; i < n; i++)
    {
        cout << ans[i] << "\n";
    }

    return 0;
}