[골드 2] 백준 2878 - 캔디캔디 (파이썬)

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

풀이

총 M개의 캔디를 N명의 학생에게 나눠줘야 한다.

각 학생은 자신이 원하는 캔디 수가 있고, 그보다 적게 받으면 “화남 지수”가 받은 개수^2로 증가한다.

M개를 어떻게 배분해야 전체 화남 지수의 합이 최소가 되도록 하는지가 문제이다.

사용한 풀이 전략은 다음과 같다.

1. 이분 탐색을 이용해 평균 컷 구하기

• 학생마다 min(A[i], mid) 만큼 주는 방식을 생각한다.
• 이 때 mid를 조절하며, 총 줄 수 있는 캔디가 M 이상이 되는 최소 mid를 찾는다.
• mid 이상을 주면 안 되고, 모든 학생에게 최대한 공평하게 적게 주자는 전략.

2. 이분 탐색 후 배분

• 최적의 컷 수 p를 찾았으면,
• 모든 학생에게 min(p - 1, A[i])만큼 먼저 배분한다.
• 남은 캔디가 있다면, A[i] > 0인 학생에게 하나씩 더 나눠준다.

3. 화남 지수 계산

• 학생별로 받은 캔디 수의 제곱을 합해 최종 정답을 구한다.

M, N = map(int, input().split())
A = [int(input()) for _ in range(N)]

M은 전체 캔디 수, N은 학생 수, A는 각 학생이 원하는 캔디 수이다.

debt = max(0, sum(A) - M)
low, high, p = 0, max(A), -1

while low <= high:
    mid = (low + high) // 2
    pay_back = sum(min(mid, x) for x in A)

    if debt <= pay_back:
        p = mid
        high = mid - 1
    else:
        low = mid + 1

pay_back은 각 학생이 mid 이상 못 받게 했을 때 줄 수 있는 총량이다.

debt는 총 원하는 양 - M, 즉 못 주는 양이다.

p를 찾는 이유는 학생들이 p이상 받으면 안되게끔 제한하기 위해서 사용된다.

캔디를 배분한다.

angry = [0] * N
for i in range(N):
    if A[i] >= p - 1:
        debt -= p - 1
        angry[i] += p - 1
        A[i] -= p - 1
    else:
        debt -= A[i]
        angry[i] += A[i]
        A[i] = 0

각 학생에게 먼저 p - 1 또는 가능한 만큼 배분한다.

if debt:
    for i in range(N):
        if debt > 0 and A[i] > 0:
            debt -= 1
            angry[i] += 1
            A[i] -= 1

그리고 남은 캔디가 있으면 하나씩 더 준다.

그리고 화남 지수를 계산해준다.

answer = 0
for i in range(N):
    answer += (angry[i] ** 2) % mod

print(answer)

각 학생의 화남 지수는 (받은 캔디 수) ** 2이고, 이를 모두 더해서 출력한다.

결론적으로, 이 문제는 단순히 그리디하게 큰 학생부터 나눠주는 것이 아닌, 캔디 개수의 상한선을 이분 탐색으로 찾고, 그 기준으로 나눠주는 공평한 분배 전략이 필요하다.

코드

# 백준 2878 - 캔디캔디
# 분류 : 이분 탐색

M, N = map(int, input().split())
A = [int(input()) for _ in range(N)]

mod = 2 ** 64

debt = max(0, sum(A) - M)

low, high, p = 0, max(A), -1

while low <= high :
    mid = (low + high) // 2

    pay_back = 0
    for i in range(N) :
        if A[i] >= mid :
            pay_back += mid
        else :
            pay_back += A[i]

    if debt <= pay_back :
        p = mid
        high = mid - 1
    else :
        low = mid + 1

# p, p - 1, A[i]

angry = [0] * N
for i in range(N) :
    if A[i] >= p - 1 :
        debt -= p - 1
        angry[i] += p - 1
        A[i] -= p - 1
    else :
        debt -= A[i]
        angry[i] += A[i]
        A[i] = 0

if debt :
    for i in range(N) :
        if debt > 0 and A[i] > 0 :
            debt -= 1
            angry[i] += 1
            A[i] -= 1

answer = 0
for i in range(N) :
    answer += (angry[i] ** 2) % mod

print(answer)