문제 설명
지훈이가 최근에 즐기는 컴퓨터 게임이 있다. 이 게임은 여러 플레이어가 참여하며, 각 플레이어는 특정한 색과 크기를 가진 자기 공 하나를 조종하여 게임에 참여한다. 각 플레이어의 목표는 자기 공보다 크기가 작고 색이 다른 공을 사로잡아 그 공의 크기만큼의 점수를 얻는 것이다. 그리고 다른 공을 사로잡은 이후에도 본인의 공의 색과 크기는 변하지 않는다. 다음 예제는 네 개의 공이 있다. 편의상 색은 숫자로 표현한다.
공 번호색크기
1 | 1 | 10 |
2 | 3 | 15 |
3 | 1 | 3 |
4 | 4 | 8 |
이 경우, 2번 공은 다른 모든 공을 사로잡을 수 있다. 반면, 1번 공은 크기가 더 큰 2번 공과 색이 같은 3번 공은 잡을 수 없으며, 단지 4번 공만 잡을 수 있다.
공들의 색과 크기가 주어졌을 때, 각 플레이어가 사로잡을 수 있는 모든 공들의 크기의 합을 출력하는 프로그램을 작성하시오.
제한 사항
풀이
문제를 요약하면, 크기와 색깔을 갖는 N개의 공이 입력으로 주어질 때, 자신보다 작고 색깔이 다른 공의 크기의 총합을 구하는 것이다.
모든 공을 살펴보며 비교하면 문제를 쉽게 풀 수 있지만, $O(N^2)$의 시간 복잡도를 갖게 된다.
해당 문제는 N의 최대 크기가 200'000이기 때문에 시간초과가 발생할 것이다.
따라서, 다른 접근이 필요하다.
누적합을 이용하는 방법이 존재한다.
만약, 공이 크기를 기준으로 정렬이 되어 있다면 자신보다 앞에 있는 공들만 영향을 주게 된다.
자신보다 뒤에 있는 공들은 자신 보다 크기 때문이다.
그렇다면, 이 부분에서 색깔에 대해서만 필터링해 주면 된다.
자신이 제외할 부분은 자신의 색깔과 같은 공이다.
따라서, 색깔별로 현재까지의 합을 저장했다면 쉽게 값을 찾아 제외할 수 있다.
한 가지 주의할 점은 만약 자신과 크기가 같은 공이 있을 수 있다.
해당 공은 자신에게 영향을 줄 수 없다.
문제의 조건을 보면 작은 공만 취급하기 때문이다.
이를 계산하는 공식은 다음과 같다.
//sum = 누적합(모든 공)
//Colors[] = color 별 누적합
//Sizes[] = size 별 누적합
Answer[idx] = sum - Colors[color] - Sizes[size];
또한, 이전 공과 크기, 색깔 모두 일치하는 경우도 존재할 수 있다.
해당 경우는 모두 같은 결과를 가져야 한다.
if (Balls[i - 1].color == color && Balls[i - 1].size == size) Answer[idx] = Answer[Balls[i - 1].idx];
전체 코드
#include <stdio.h>
#include <string>
#include <vector>
#include <algorithm>
#include <iostream>
#include <sstream>
#include <cmath>
#include <climits>
#include <queue>
#include <map>
#include <unordered_map>
#include <set>
using namespace std;
const int MAX = 200'020;
int N;
struct ball {
int size, color, idx;
};
bool cmp(ball& a, ball& 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;
vector<ball> Balls;
vector<int> Colors(MAX, 0);
vector<int> Sizes(MAX, 0);
vector<int> Answer(MAX, 0);
for (int i = 0; i < N; i++)
{
int size, color;
cin >> color >> size;
Balls.push_back({ size,color,i });
}
sort(Balls.begin(), Balls.end(), cmp);
int sum = 0;
for (int i = 0; i < N; i++)
{
int size = Balls[i].size;
int color = Balls[i].color;
int idx = Balls[i].idx;
Answer[idx] = sum - Colors[color] - Sizes[size];
Colors[color] += size;
Sizes[size] += size;
sum += size;
if (i == 0) continue;
if (Balls[i - 1].color == color && Balls[i - 1].size == size) Answer[idx] = Answer[Balls[i - 1].idx];
}
for (int i = 0; i < N; i++)
{
cout << Answer[i] << "\n";
}
return 0;
}