컬러볼
문제 해석
각 플레이어의 공은 색깔과 크기가 숫자로 표현되어있습니다.
자기 공보다 크기가 작고 색이 다른 공을 사로잡아 그 공의 크기만큼 점수를 얻습니다.
본인의 공 색깔과 크기는 변형되지 않습니다.
각 플레이어가 사로잡을 수 있는 모든 공들의 크기의 합을 구하는 문제입니다.
설계(20)
Naive하게 한 플레이어가 모든 공을 조사하는데 O(N) x 모든 플레이어가 이를 반복 O(N) = O(N^2)으로 시간초과가 납니다.
정렬을 이용하면 이 문제를 O(NlogN)에 풀 수 있습니다.
우선, {공의 번호, 색, 크기}로 구성된 자료구조를 크기순으로 오름차순 정렬을 합니다.(O(NlogN))
그리고 순차적으로 총 누적합을 구하고, 각 색깔마다 누적합을 카운트해서 각 플레이어가 얻을 수 있는 점수를 구합니다.
여기서 주의할 점은, 크기가 같은 경우입니다. 위의 로직대로 진행하면 색은 다르고 크기가 같은 공에 대해서도 카운트 되는 오류가 발생할 것입니다.
이를 처리하기 위하여 이전의 공 색깔,공 크기를 이용하여 공색깔이 같고 공 크기가 같은 경우와 공 색깔이 다르고 공 크기가 같은 경우, 이외의 경우로 나누어서 처리를 해주어야 합니다.
첫번째 방법은 구현은 가능하지만 크기가 같은 경우를 고려해야하는 어려움이 있었습니다. 다른 풀이로는 해당 블로그를 참조하였을 때, O(N)에 풀이가 가능하고 쉽게 구현할 수 있습니다.
공의 크기가 <= 2000인 것을 감안하여 2000개 짜리 벡터를 생성합니다. 그 후, 각 크기에 대해서 첫번 째 풀이의 로직과 같이 구현하면 공의 크기가 같은 경우는 고려를 안해도 되기 때문에 쉽게 구현할 수 있습니다.
구현(40)
코드1(O(NlogN))
#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;
#define MAXN 200000
int N;
struct node {
int idx;
int color;
int size;
bool operator < (const node& a)const {
if (size == a.size) return color < a.color;
return size < a.size;
}
};
node arr[MAXN];
int ans[MAXN];
int cnt[MAXN + 1];
int main() {
ios::sync_with_stdio(false);
cin.tie(0);
cin >> N;
for (int i = 0; i < N; i++) {
cin >> arr[i].color >> arr[i].size;
arr[i].idx = i;
}
sort(arr, arr + N);
int total = 0;
int prev_total = 0;
int prev_size = 0;
int prev_color = 0;
int tmp = 0;
for (int i = 0; i < N; i++) {
if (arr[i].size == prev_size) {
if (prev_color != arr[i].color) {
prev_total += tmp;
ans[arr[i].idx] = total - cnt[arr[i].color] - prev_total;
tmp = arr[i].size;
}
else {
tmp += arr[i].size;
ans[arr[i].idx] = total - cnt[arr[i].color] - prev_total;
}
}
else {
ans[arr[i].idx] = total - cnt[arr[i].color];
prev_total = 0;
tmp = arr[i].size;
}
prev_color = arr[i].color;
prev_size = arr[i].size;
cnt[arr[i].color] += arr[i].size;
total += arr[i].size;
}
for (int i = 0; i < N; i++) {
cout << ans[i] << "\n";
}
return 0;
}
코드2 C++(O(N))
#include <iostream>
#include <vector>
using namespace std;
#define MAXN 200000
int ans[MAXN];
int cnt[MAXN + 1];
vector<pair<int, int> > adj[2001];
int N;
int main() {
ios::sync_with_stdio(false);
cin.tie(0);
cin >> N;
for (int i = 0; i < N; i++) {
int color, size;
cin >> color >> size;
adj[size].push_back({ i,color });
}
int total = 0;
for (int i = 1; i <= 2000; i++) {
for (int j = 0; j < adj[i].size(); j++) {
ans[adj[i][j].first] = total - cnt[adj[i][j].second];
}
for (int j = 0; j < adj[i].size(); j++) {
cnt[adj[i][j].second] += i;
total += i;
}
}
for (int i = 0; i < N; i++) {
cout << ans[i] << "\n";
}
return 0;
}
코드3 Python(O(N))
N = int(input())
arr = [list(map(int,input().split())) for i in range(N)]
ans = [0]*N
cnt = [0]*200001
total = 0
adj = [[] for _ in range(2001)]
for i in range(N):
size = arr[i][1]
color = arr[i][0]
adj[size].append([i,color])
for i in range(1,2001):
for idx,color in adj[i]:
ans[idx] = total - cnt[color]
for _,color in adj[i]:
cnt[color] += i
total += i
for i in ans:
print(i)
디버깅
제출결과
c++ 88 ms(O(NlogN)) c++ 72 ms(O(N)) pypy3 792 ms (O(N))
마무리
생각보다 아이디어를 생각하기 어려웠던 문제였습니다. 여러 방법으로 풀이가 가능하여 공부하기에 좋은 문제인 것 같습니다.