[Python 3] BOJ - 2878 "캔디캔디"

 # 문제 링크

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

2878번: 캔디캔디

 

 # 풀이

 개인적으로 이 문제의 핵심은 이분 탐색이라고 생각한다. 여러 포스팅에서 보았듯이, 이분 탐색의 핵심은 특정 x에 대한 대답이 YES일 때, x+1, x+2, ... 에 대한 대답도 YES가 되는지를 파악하는 것이다. 그래야 반씩 줄여서 탐색하는 것이 의미가 있기 때문이다

 따라서 우리가 최소화 하고 싶은 것은 각 불만의 제곱의 합이다. 그런데 그냥 더하는 것이 아니라 제곱한 후 더하기 때문에 최대값에 영향을 굉장히 크게 받는다. 따라서, 우리는 원하는 사탕값의 최대값을 최소화해야한다. 따라서, 현재 원하는 사탕의 양을 균일하게 맞춰준다면 최대값을 일괄적으로 조절할 수 있다.

<그림 1. 그리디 증명>

 위 그림 1에서 오른쪽은 특정 k번째 값으로 최대값을 고르게 만든 상황이다. 왼쪽은 아무 사탕도 주지 않은 초기 상황이고 이를 내림차순으로 정렬한 결과이다. 이 때 k번째 이후로는 최대값을 고르게 만들 필요가 없으므로 사실상 a와 값이 같아진다. 그런데 오른쪽 그림과 같이 만들 수가 있다면 결국 나누어 준 사탕은 m보다 작거나 같은 값이 됨을 알 수 있다. 

 따라서, 수식으로 표현하면 $(a_1 - a_k) + (a_2 - a_k) + ... + (a_{k-1} - a_k) \leq m$이 된다. 그런데 그리디하게 맨 앞에서 k까지 다 주려면 필요한 사탕의 양은 $\sum_{i = 1}^{k} a_i = q$가 된다. 따라서, m에서 q를 빼면 $-ka_k$이고 값은 음수가 됨을 알 수 있다. 역시 k - 1까지 준다고 해도, 음수가 나오기 때문에 최소한 두 개 이상의 값이 그대로 a의 값과 같아진다.

<그림 2. 그리디 증명>

 위 그림은 증명 결과의 시각화이다. 결국엔 p부터 k까지 적어도 2개 이상의 더 큰 값이 존재하고, 각각은 제곱만큼 기여하기 때문에 나머지 $$b_1, b_2, ... b_p$$까지의 기여분보다 훨씬 큼을 알 수 있다.

 그런데 이분 탐색을 하면 끝일까? 아니다. 이분 탐색을 해도 줄 수 있는 사탕값이 남아 있을 수 있다. 그런데 이분 탐색 후에 남은 사탕값은 절대로 모든 사람들에게 나눠줄 수 없다. 그런데 앞에서 최대값을 줄이는 것이 최적이므로, 즉 불만값이 고르면 고를수록 값이 큰 아웃라이어에 대한 영향을 최소화할 수 있기 때문에, 남은 불만값도 줄 수 있는 사탕값에 따라 최대한 고르게 나눠주면 시간안에 정답을 구할 수 있다.

 

 # 코드

import sys, heapq

# check : 특정 높이 num만큼 고르게 할 수 있는지 체크하는 함수
def check(num):
    cnt = m
    for i in arr:
        if i > num:
            cnt -= (i - num)
    return cnt

# 입력부
m, n = map(int, sys.stdin.readline().split())
arr = []
for _ in range(n):
    arr.append(int(sys.stdin.readline()))
    
# 이분탐색
left = 0
right = 9876543210
while left < right:
    mid = (left + right) // 2
    if check(mid) >= 0:
        right = mid
    else:
        left = mid + 1

# 이분 탐색에서 구한 최대 높이만큼 사탕 분배
for i in range(n):
    if arr[i] <= right:
        continue
    else:
        m -= abs(right - arr[i])
        # 최대 높이보다 더 크면 갱신
        arr[i] = right

# 최대 힙에 arr값들을 전부 추가
q = []
for i in arr:
    heapq.heappush(q, -i)

# 남은 사탕 갯수를 고르게 분배
while m:
    m -= 1
    now = heapq.heappop(q)
    # 왜 now + 1인가?
    # 왜냐하면 최대 힙을 구현하기 위해 음수값을 넣었기 때문
    heapq.heappush(q, now + 1)

# 정답 출력
ans = 0
for i in q:
    ans += i ** 2
print(ans % pow(2, 64))