[백준]JAVA - 10800번: 컬러볼

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

10800번: 컬러볼

풀이

공의 색깔과 크기를 비교해서 i번째 공으로 잡아먹을 수 있는 공들의 크기 합을 구하는 문제이다.
완전 탐색으로 접근하면 공의 개수가 200,000까지의 크기를 가지므로 시간 초과가 발생할 것이다.
따라서 위 문제는 누적합과 정렬을 이용해서 해결해야한다.

누적합을 사용하기 위해서는 공의 크기가 작은 순으로 정렬을 해주어야한다.

public static class Ball implements Comparable<Ball> {
    int idx;
    int c;
    int s;
		
    public Ball (int idx, int c, int s) {
        this.idx = idx;
        this.c = c;
        this.s = s;
    }
		
    public int compareTo (Ball o) {
        return this.s - o.s;
    }
}

정렬을 재정의하는 방법은 크게 두 가지이다. 블로그에서 자주 언급했으므로 자세한 건 정렬 카테고리에 들어가서 확인해보면 된다. 알고리즘에서 자주 등장하므로 꼭 알아두는 게 좋다.

정렬을 완료했으면 누적합만 구하면 된다.
여기서 왜 누적합을 이용하는지에 대해서 궁금할 수도 있다.

위의 그림과 같이 색깔이 다른 4개의 공을 크기 순으로 정렬했다.
문제에서는 "색깔이 같은 건 먹을 수 없다." 조건이 있기 때문에 누적합과 색깔 별 누적합을 가지고 있어야한다.

idx	전체 누적합	색깔 별 누적합
0 (빨간색)	1	빨간색: 1, 초록색: 0, 파란색: 0
1 (초록색)	3	빨간색: 1, 초록색: 2, 파란색: 0
2 (파란색)	7	빨간색: 1, 초록색: 2, 파란색: 4
3 (빨간색)	14	빨간색: 8, 초록색: 2, 파란색: 4

위와 같이 누적합을 구해주면서 i번째 도달했을 때 전체 누적합에서 i번째 공의 색깔 누적합을 빼주면 i번째 공이 먹은 합이 나오는 것이다.

for (int i=0; i<N; i++) {
    Ball a = ball.get(i);
		
    sum += a.s;
    color[a.c] += a.s;    
			
    ans[a.idx] = sum - color[a.c];
}

위에 말한대로 코드를 짜게 되면 대충 이런식으로 코드가 나올 것이다. 이건 올바른 코드가 아니다. 왜냐하면 같은 크기의 공을 먹을 수 없다는 걸 간과했다.

따라서 i번째의 공 크기가 i-1번째 공의 크기보다 확실하게 클 때 합을 누적해야한다.

int sum = 0;
int j = 0;
for (int i=0; i<N; i++) {
    Ball a = ball.get(i);
    Ball b = ball.get(j);
			
    while (b.s < a.s) {
        sum += b.s;
        color[b.c] += b.s;
				
        b = ball.get(++j);
    }
			
    ans[a.idx] = sum - color[a.c];
}

<최종코드>

import java.io.*;
import java.util.*;

public class Main {
    public static class Ball implements Comparable<Ball> {
        int idx;
        int c;
        int s;
		
        public Ball (int idx, int c, int s) {
            this.idx = idx;
            this.c = c;
            this.s = s;
        }
		
        public int compareTo (Ball o) {
            return this.s - o.s;
        }
    }
	
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st;
		
        int N = Integer.parseInt(br.readLine()); // 공의 개수
		
        ArrayList<Ball> ball = new ArrayList<>();
		
        for (int i=0; i<N; i++) {
            st = new StringTokenizer(br.readLine());
			
            int color = Integer.parseInt(st.nextToken());
            int size = Integer.parseInt(st.nextToken());
			
            ball.add(new Ball(i, color, size));
        }
		
        Collections.sort(ball);
		
        int[] color = new int[N+1];
        int[] ans = new int[N];
		
        int sum = 0;
        int j = 0;
        for (int i=0; i<N; i++) {
            Ball a = ball.get(i);
            Ball b = ball.get(j);
			
            while (b.s < a.s) {
                sum += b.s;
                color[b.c] += b.s;
				
                b = ball.get(++j);
            }
			
            ans[a.idx] = sum - color[a.c];
        }
		
		
        for (int num : ans) {
            System.out.println(num);
        }
    }
}