[백준] 10800번: 컬러볼

📌 문제

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

📌 설명

  공이 다른 공을 사로잡을 수 있다. 사로잡기 위한 조건은 색이 다르거나, 크기가 더 작을때이다. 공에 대한 정보를 배열에 저장하고 브루트포스로 일일히 계산하게 되면 $O(N^2)$로 시간 초과가 발생하게 된다. 정렬과 누적 합 개념을 사용해야 한다.

 공에 대한 정보를 저장하는 구조체 Ball을 선언한다. 공의 색, 크기, 인덱스 정보를 기록한다. 그 다음 공을 크기 순으로 오름차순 정렬한다. 크기가 같다면 색을 기준으로 정렬한다. 중복된 공이 들어올 수 있으며, 이에 대한 처리를 효율적으로 하기 위함이다.

 정렬된 배열에 접근하면서 공이 사로잡을 수 있는 크기를 효율적으로 구한다. 이때 누적 합 개념이 사용된다.

📌 코드

#include <iostream>
#include <algorithm>

using namespace std;

struct Ball {
	int color;
	int size;
	int idx;
};

Ball balls[200001];
int cumsumByColor[200001];
int cumsumBySize[2001];

int ans[200001];

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

int main() {
	int n;
	cin >> n;

	for (int i = 0; i < n; i++) {
		cin >> balls[i].color >> balls[i].size;
		balls[i].idx = i;
	}

	sort(balls, balls + n, cmp);

	// i번 째 공이 사로잡은 공의 크기: i번 째까지 모든 공의 크기의 합 
	// - i번 째 공과 색이 같은 공의 크기의 합(1) - i번 째 공과 크기가 같은 공의 크기의 합(2)
	// + i번 째 공의 크기(3) ((2)에서 불필요하게 빠진 무게 보정)
	int total = 0;
	for (int i = 0; i < n; i++) {
		int color = balls[i].color;
		int size = balls[i].size;
		int idx = balls[i].idx;

		total += size;
		cumsumByColor[color] += size;
		cumsumBySize[size] += size;

		// 현재 공이 중복된 공인 경우 이전 공의 정보를 가져온다.
		if (i - 1 >= 0 && balls[i - 1].color == color && balls[i - 1].size == size) {
			ans[idx] = ans[balls[i - 1].idx];
		}
		else
			ans[idx] = total - cumsumByColor[color] - cumsumBySize[size] + size;
	}

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

---

	// i번 째 공이 사로잡은 공의 크기: i번 째까지 모든 공의 크기의 합 
	// - i번 째 공과 색이 같은 공의 크기의 합(1) - i번 째 공과 크기가 같은 공의 크기의 합(2)
	// + i번 째 공의 크기(3) ((2)에서 불필요하게 빠진 무게 보정)
	int total = 0;
	for (int i = 0; i < n; i++) {
		int color = balls[i].color;
		int size = balls[i].size;
		int idx = balls[i].idx;

		total += size;
		cumsumByColor[color] += size;
		cumsumBySize[size] += size;

		// 현재 공이 중복된 공인 경우 이전 공의 정보를 가져온다.
		if (i - 1 >= 0 && balls[i - 1].color == color && balls[i - 1].size == size) {
			ans[idx] = ans[balls[i - 1].idx];
		}
		else
			ans[idx] = total - cumsumByColor[color] - cumsumBySize[size] + size;
	}

 현재 i번 째 공이 사로잡을 수 있는 크기 ans[idx]를 계산하고자 한다(idx: i번 째 공의 실제 인덱스) . 세 가지 누적 합 변수를 사용한다. total은 0~i번 째까지 공의 크기의 합이다. cumsumByColor[c]는 색깔이 c인 공의 크기의 합이다. cumsumBySize[s]는 크기가 s인 공의 크기의 합이다. 식은 다음과 같다. ans[idx] = total - cumsumByColor[c] - cumsumBySize[s] + s(c: 현재 공의 색상, s: 현재 공의 크기). 마지막에 s를 더하는 이유는 cumsumBySize[s]를 뺄 때 불필요하게 현재 공의 크기 또한 빠졌기 때문이다. 이를 보정하기 위해 s를 더하는 것이다.