[백준 - 2878] 캔디캔디 - Java

문제 링크

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

문제

티어: 골드 2
알고리즘: 그리디, 정렬, 누적합?

입력

첫째 줄에 M(1 ≤ M ≤ 2×109)와 N(1 ≤ N ≤ 100,000)이 주어진다. 둘째 줄부터 N개의 줄에는 각 친구들이 받고 싶어하는 사탕의 개수가 주어진다. 이 개수는 2×109보다 작으며, 친구들이 받고 싶어하는 사탕의 개수의 합은 항상 M을 넘는다.

출력

첫째 줄에 택희 친구들의 분노의 합의 최솟값을 264로 나눈 나머지를 출력한다.

풀이

구하고자 하는 것은 친구들의 분노의 합을 최소화한 값이다.

아이디어는 간단하다. 분노의 합은 남은 사탕의 제곱이기 때문에 현재 친구들 중 남은 사탕이 큰 친구에게 우선적으로 사탕을 나눠주는 것이 최선의 선택이다.

나는 이를 구현하기 위해 정렬과 누적합을 사용했다. (누적합은 꼭 필요한진 모르겠지만, 나는 필요할 것 같아 사용했다.)

예제를 이용해 설명하겠다.

M은 10, N은 4이고,
각 친구들은 4, 5, 2, 3을 요구한다.

먼저, 내림차순으로 정렬한다. -> 5 4 3 2

그러면 가장 먼저 사탕을 줘야될 사람은 자연스럽게 낮은 인덱스 순이 된다.

그래서 0번 째 인덱스에 사탕을 주는데, 1번 째 인덱스 사탕의 개수와 맞춰준다.

그러면 4 4 3 2이고, M은 9가 된다. 이때 나는 sum이라는 배열을 둬서 0번 째 인덱스에 -1를 해줬다. 그래서 sum은 -1 0 0 0이 된다.

다음은 0 번째와 1 번째 인덱스를 3으로 맞춰준다.

그러면 3 3 3 2가 되고, M은 7이 된다. sum은 -1 -1 0 0이 된다.

다음은 0, 1, 2 번째 인덱스를 2로 맞춰준다.

그러면 2 2 2 2가 되고, M은 4가 된다. sum은 -1 -1 -1 0이 된다.

이제 모든 값이 같아졌기 때문에 앞에서부터 차례대로 사탕을 1씩 주면된다.

그러면 1 1 1 1이 되고, M은 0이 되며, sum은 -1 -1 -1 -1이 된다.

sum의 누적합을 구하면 -4 -3 -2 -1이 되고, 초기의 값인 5 4 3 2를 인덱스의 맞게 계산해주면, 남은 사탕의 개수가 된다.

그래서 최종 개수는 1 1 1 1로, 최솟값은 4가 된다.

이 풀이의 핵심은 남은 사탕의 개수가 많은 사람들에게 최대한 골고루 나눠주는 것이다. 이는 제곱의 합의 최솟값을 보장한다.

또한 누적합과 나눠줄 사탕의 개수는 식으로 한번에 구할 수 있기 때문에 O(N) 풀이가 가능하다.

소스 코드

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

public class Main {
    static int M, N;
    // static int MOD = Math.pow(10, )
    public static void main(String args[]) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        M = Integer.parseInt(st.nextToken());
        N = Integer.parseInt(st.nextToken());
        Integer[] list = new Integer[N];
        int[] sum = new int[N];

        for (int i = 0; i < N; i++) {
            list[i] = Integer.parseInt(br.readLine());
        }

        Arrays.sort(list, Collections.reverseOrder());

        int cursor = 0;
        while (M != 0) {
            cursor += 1;
            if (cursor == N) {
                sum[cursor - 1] -= M / cursor;
                if (M % cursor != 0) {
                    sum[(M % cursor) - 1] -= 1;
                }
                M = 0;
            } else {
                int diff = list[cursor - 1] - list[cursor];

                if (diff * cursor <= M) {
                    M -= diff * cursor;
                    sum[cursor - 1] -= diff;
                } else {
                    sum[cursor - 1] -= M / cursor;
                    if (M % cursor != 0) {
                        sum[(M % cursor) - 1] -= 1;
                    }
                    M = 0;
                }
            }
        }
        
        for(int i=N - 1; i >= 1; i--) {
            sum[i - 1] += sum[i];
        }
        
        BigInteger answer = BigInteger.ZERO;
        BigInteger MOD = BigInteger.ONE.shiftLeft(64); // 2^64
        for(int i=0; i<N; i++) {
            int r = list[i] + sum[i];
            BigInteger remain = BigInteger.valueOf(r);
            answer = answer.add(remain.multiply(remain)).mod(MOD);
        }
        System.out.println(answer);
    }
}