백준(Baekjoon)문제 풀이 - 10800 컬러볼(C++)

컬러볼

문제

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

같은 색의 공끼리는 먹을 수 없고, 자기보다 크기가 작은 공끼리만 먹을 수 있을 때, 주어진 공이 공을 얼마나 먹을 수 있는지, 먹을 수 있는 공의 총 합을 출력하면 되는 문제이다.

입출력 형식

<입력 형식>

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

<출력 형식>

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

 

<입력 예>

4
1 10
3 15
1 3
4 8

 

<출력 예>

8
21
0
3

아이디어

 N 범위가 크기 때문에 O(N^2)은 당연히 시간 초과가 난다. 이 문제를 보고 처음에는 Dp인줄 알았는데, 누적합 + 정렬 문제였다. 

 

Q1 우선, 같은 색의 공끼리 먹을 수 없다는 조건이 없다면 어떻게 문제를 풀 수 있는지 생각해 보자.

 

-> 생각보다 간단하다. 정렬한 후, sum을 더해가며 자신보다 크기가작은 모든 공을 더한 뒤, 크기가 같은 공을 빼주면 될 것이다.(정렬을 한다 해도 같은 수들까지 걸러지진 않음)

 

Q2 그렇다면, 여기에서 같은 공끼리는 먹을 수 없다는 조건이 추가된다면 어떨까?

 

-> 이러한 경우에는 생각하기 까다로울 수 있지만, 누적합 배열을 하나 만드는 것으로 해결할 수 있다. <sum_col[i] = i 색으로 칠해진 공의 누적합 배열> 을 하나 만들어주고, 전체 누적 합에서 이를 빼준다면, 같은 색을 제외하고 먹을 수 있는 모든 공을 먹은 값이 나오게 된다.

 

-> 위 두가지 사실을 종합해본다면, 전체 누적합에서 같은 색을 갖는 누적합과, 같은 수를 갖는 누적합을 빼준 뒤, 교집합인 현재의 크기값을 더해주면 답이 된다.

 

※만약 전의 값과 색이 현재의 값과 색과 모두 같다면, 위 계산을 똑같이 실행하면 틀리기 때문에 예외처리를 반드시 해주어야 한다.

※정렬을 위한 compare함수 작성시 크기를 기준으로 하되 두 크기가 같다면, 색을 2번째 우선순위로 반드시 두어야 한다.

 

<소스 코드>

#include <stdio.h>
#include <algorithm>
#include <vector>

using namespace std;

long long n;

struct dat
{
	long long idx, col, val;
};

bool comp(dat a, dat b)
{
	if (a.val == b.val) return a.col < b.col;
	return a.val < b.val;
}

long long sum_col[200010], sum_val[200010];
long long ans[200010];
vector<dat> vec;

int main()
{
	long long i, j;
	scanf("%lld", &n);
	for (i = 0; i < n; i++) {
		long long x, y;
		scanf("%lld %lld", &x, &y);
		vec.push_back({ i, x, y });
	}
	sort(vec.begin(), vec.end(), comp);
	long long sum = 0;
	for (i = 0; i < n; i++) {
		long long col = vec[i].col, val = vec[i].val;
		long long idx = vec[i].idx;
		sum += val;
		sum_col[col] += val;
		sum_val[val] += val;
		if (i > 0 && val == vec[i - 1].val && col == vec[i - 1].col) ans[idx] = ans[vec[i - 1].idx];
		else ans[idx] = sum - sum_col[col] - sum_val[val] + val;
	}
	for (i = 0; i < n; i++) printf("%lld\n", ans[i]);
	return 0;
}

아이디어를 떠올리기는 생각보다 쉬웠으나, 예외처리가 빡셌다.