[알고리즘 / 백준] 10800 - 컬러볼

www.acmicpc.net/problem/10800

10800번: 컬러볼

 

문제는 위와 같으며, 현재 공의 크기보다 작은 공을 잡을 수 있다는 것이 중요한 조건이기 때문에 우선 공을 크기 순으로 정렬한 리스트로 저장하였습니다.

 

그리고 앞에서부터 공을 뽑고 그 합을 알고 있으면, 뒤에 오는 공은 항상 앞에 뽑은 공보다 크기가 크기 때문에 앞서 뽑은 모든 공의 크기에서 현재 공의 색과 같은 공들의 크기만 제거하면 잡을 수 있는 공의 크기 합을 알 수 있습니다.

 

이 부분만 자바 코드로 살펴보면 다음과 같습니다.

int[] answer = new int[n];

int sum = 0;
int[] color = new int[n + 1];

for (int i = 0; i < n; i++) {
    Ball a = balls.get(i);
    
    color[a.color] += a.size;
    sum += a.size;
  
    answer[a.idx] = sum - color[a.color];
}

 

하지만 이렇게 하는 경우, 크기가 같은 경우를 처리할 수 있는 방법이 없습니다. 따라서 공의 크기를 누적할 때 확실하게 크기가 작은 공의 크기만을 누적할 수 있도록 변수를 하나 더 두어 처리하였습니다.

int[] answer = new int[n];

int sum = 0;
int[] color = new int[n + 1];

for (int i = 0, j = 0; i < n; i++) {
    Ball a = balls.get(i);
    Ball b = balls.get(j);

    while (b.size < a.size) {
        sum += b.size;
        color[b.color] += b.size;

        b = balls.get(++j);
    }

    answer[a.idx] = sum - color[a.color];
}

---

전체 코드를 살펴보면 다음과 같습니다.

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.*;

public class Main {
    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int n = Integer.parseInt(br.readLine());

        ArrayList<Ball> balls = new ArrayList<>();

        for (int i = 0; i < n; i++) {
            StringTokenizer st = new StringTokenizer(br.readLine());
            int c = Integer.parseInt(st.nextToken());
            int s = Integer.parseInt(st.nextToken());
            balls.add(new Ball(i, c, s));
        }

        // 크기가 작은 순으로 Ball 리스트 정렬
        Collections.sort(balls, new Comparator<Ball>() {
            @Override
            public int compare(Ball o1, Ball o2) {
                return o1.size - o2.size;
            }
        });

        int[] answer = new int[n];

        int sum = 0;
        int[] color = new int[n + 1];
        for (int i = 0, j = 0; i < n; i++) {
            Ball a = balls.get(i);
            Ball b = balls.get(j);

            // 이전 공의 크기가 현재 공의 크기보다 작은 경우만
            while (b.size < a.size) {
                sum += b.size;  // 작은 공의 크기 누적
                color[b.color] += b.size;  // 작은 공의 색 인덱스 위치에 작은 공의 크기를 추가

                b = balls.get(++j);  // 다음 공 뽑기
            }

            // 현재 공이 잡을 수 있는 공의 수 (크기가 작은 전체 누적 공 크기 - 현재 공과 같은 색의 공 크기 합)
            answer[a.idx] = sum - color[a.color];
        }

        // 결과 출력
        StringBuilder sb = new StringBuilder();
        for (int a : answer) {
            sb.append(a).append("\n");
        }

        System.out.println(sb);

    }

    // Ball 클래스 (인덱스, 색, 크기)
    private static class Ball {
        int idx;
        int color;
        int size;

        public Ball(int idx, int color, int size) {
            this.idx = idx;
            this.color = color;
            this.size = size;
        }
    }
}

 

파이썬 코드는 다음과 같습니다.

from sys import stdin


def return_size(ball):
    return ball[2]


n = int(stdin.readline())

balls = []
for i in range(n):
    c, s = map(int, stdin.readline().split())
    balls.append([i, c, s])

balls.sort(key=return_size)

s = 0
j = 0
color = [0] * (n + 1)
answer = [0] * n
for i in range(n):
    a = balls[i]
    b = balls[j]

    while b[2] < a[2]:
        s += b[2]
        color[b[1]] += b[2]
        j += 1
        b = balls[j]

    answer[a[0]] = s - color[a[1]]

print("\n".join(map(str, answer)))