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

중급 알고리즘 - 그리디 수강 후 연습 문제로 풀게된 문제이다.

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

2878번: 캔디캔디

친구들에게 사탕을 나눠줘야하는데 사탕은 반드시 모자라다.

내가 원하는 사탕이 10개인데, 사탕을 5개만 받았다면

그 친구는 못받은 사탕개수의 제곱만큼 분노하게 된다.

(분노치 5*5 = 25)

 

각 친구들의 분노치 합을 최소로 만드는 문제이다.

 

이제부터 내가 이 문제를 풀기까지 사고한 과정을 정리하고자 한다.

(미래의 나에게 도움이 분명 될 것이라고 생각한다 ㅋㅋ)

---

내가 처음으로 떠올린 아이디어는 다음과 같다.

"사탕을 많이 받고싶어하는 사람에게 먼저 분배하기"

 

제곱의 합의 값을 최소로 하기 위해서는 가장 큰 값부터 없애야

가장 큰 제곱값이 점점 작아지니 최소에 가까워진다.

 

직관적이면서도 괜찮은 아이디어같다.

예제 입력에 대해 실험해보자.

3개를 원하는 친구

2개를 원하는 친구

1개를 원하는 친구

 

갖고 있는 사탕은 5개

먼저 3개를 원하는 친구에게 3개를 준다.

그 다음 2개를 원하는 친구에게 2개를 준다.

그러면 1개가 남는다! 정답과 똑같다.

(사실 이렇게 풀면 틀린다.)

 

나쁘지 않은 아이디어 같으니 그 다음으로 특이케이스를 생각해보았다.

갖고 싶은 개수가 똑같은 경우에는 어떻게 처리할까?

 

2개를 원하는 친구1

2개를 원하는 친구2

2개를 원하는 친구3

 

갖고 있는 사탕은 4개

 

먼저 친구 1에게 2개를주고

친구 2에게 2개를 주면

2개를 원하는 사람 한명이 남으므로 분노치는 4

 

하지만 이 값은 최소값이 아니다.

친구 1,2,3에게 공평하게 하나씩 나눠주고

남은 사탕 한개를 아무에게나 하나 주었을 때

1개를 원하는 사람 2명이 남으므로 분노치가 2일 때 최소값이 된다.

 

이렇게 생각해보면 이미 정답에 가까워졌다.

이 문제를 푸는 핵심은 "공평하게" 이다.

최대한 균등하게 나눠주었을 때 최소가 된다.

 

정확히는 각자 원하는 사탕의 양이 최대한 '균등하게' 되었을 때

구하는 값의 최소가 나온다.

물론 저 원하는 사탕의 양이 적은 사람의 것을 굳이 강제로 늘릴 필요는 없다.

(애초에 늘릴 수도 없다)

원하는 사탕의 양이 많은 사람부터 균등하게 줄여나가면 된다.

---

아이디어는 떠올렸다.

그러나 그 다음 문제는 구현이다.

이 아이디어를 어떻게 구현해야 할까?

 

우선 첫번째 작업은 확실하다.

1. 원하는 사탕의 양이 많은 순서대로 정렬하는 것.

 

그런데 원하는 사탕의 양이 같은 경우 이걸 중복해서 써야한다.

그러면 문제점이 발생한다.

9 9 9 9 를 균등하게 하나씩 줄이기 위해 4번의 뺄셈을 거쳐 8 8 8 8로 만든다?

한 두번이야 괜찮겠지만 이 연산을 계속 반복한다면, 또 개수가 커진다면 너무 비효율적이다.

 

그래서 2. (원하는 사탕의 개수 : 원하는 사람 수) 쌍의 자료구조를 떠올렸다.

C++기준으로는 map / 파이썬 기준으로는 딕셔너리 구조가 될 것이다.

원하는 사탕의 개수를 Key로, 원하는 사람의 수를 Value로 한다.

이후에는 위처럼 부르지 않고 단순히 키, 밸류로 지칭할 것이다.

여기에 앞서 적은 '첫번째 작업'을 섞어주면, 딕셔너리를 채우고, 키 리스트만 따로 정렬해서 쓰면 된다.

 

메모리 용량을 체크해보면 모든 자료가 중복되지 않고 하나씩 들어가더라도

최대 10만개를 저장하게 되므로 메모리 초과는 걱정하지 않아도 된다.

오히려 자료가 중복될 때 메모리 효율은 증가한다.

 

구체적인 구현방법으로 처음으로 떠올린 것은

3. 현재의 키와 그 다음의 키를 비교하여

현재 키의 값을 그 다음의 키의 값으로 맞추는 것이었다.

 

예시와 함께 설명하면 다음과 같다.

가령 현재의 키가 8이고 (8개를 원하는 사람이 존재)

그 다음 키가 4라고 해보자 (그 다음으로 4개를 원하는 사람이 존재)

 

그럼 8개를 원하는 사람에게 4개를 주어서 4개를 원하도록 만드는 것이다.

그러면 8의 키에 있는 값이 4의 키에 더해지면 된다.

 

그러나 여기에서 한가지 문제를 떠올렸다.

만약 갖고 있는 사탕이 4개보다 적어서 4의 키에 맞출 수 없다면 어떻게 해야할까?

여기에서 막히자 더 큰 문제점이 떠올랐다.

첫 키가 10억이고 그 다음 키가 1인데, 내가 갖고 있는 사탕이 1개라면 어떻게 해야할까?

그 다음키 1에 맞출 수가 없다.

 

그래서 생각을 바꿔보았다.

바로 키를 1씩 줄여보는 것이었다.

 

첫 키가 10이라면 그보다 1 작은 9에 맞춘다.

9키가 존재하면 기존 키에 더하고, 없다면 새로 만들고

 

9키에 존재하는 값들에 대해서 다시 1작은 8에 맞춘다.

8키가 존재하면 기존 키에 더하고, 없다면 새로 만들고

 

(이하 반복)

 

그러나 이 문제는 더 심각한 문제를 갖고 있었다.

내가 갖고 있을 수 있는 사탕 개수의 범위가 10억이다.

만약 내가 5억개를 갖고 있는데,

첫 키가 10억이고 그 다음키가 1이라면?

10억부터 5억까지 5억번 반복하면서 키를 왔다갔다 할 것이다.

시간초과를 피할수도 없거니와, 10억개의 키를 모두 저장하는건 비효율적이다.

(사실 pop을 쓰면 키를 최대 10만개만 저장하면서도 넣었다 뺐다를 할 수 있다.

그래도 시간초과는 못피한다)

 

그래서 다시 원래 생각으로 돌아왔다.

만약 갖고 있는 사탕이 그 다음키에 맞출 만큼 충분하지 않다면 어떻게 해야할까?

만약 갖고 있는 사탕이 그 다음키에 맞출 만큼 충분하다면 어떻게 해야할까?

 

다시 사고를 정리해보니 답이 나왔다.

 

- 만약 갖고 있는 사탕이 그 다음키에 맞출 만큼 충분하다면

=> 가진 사탕 개수를 줄이고, 그 키에 맞추면 된다.

 

- 만약 갖고 있는 사탕이 그 다음키에 맞출 만큼 충분하지 않다면

=> 공평하게 줄 수 있는 만큼 주고 나머지를 한개씩 줄 수 있을 때까지 주면 된다.

 

예시를 들어보자

내가 갖고 있는 사탕이 10개이다.

4개를 원하는 사람 2명 2개를 원하는 사람 2명이 있다.

 

데이터 저장은 {4 : 2, 2 : 2} 이다.

키를 내림차순 정렬해보자. [4, 2]

첫번째로 4를 2에 맞춘다.

 

4개를 원하는 2명에게 각각 2개씩 주어 2개를 원하도록 고치려면

2개씩 2개, 총 4개가 필요하다.

 

일반화 해보자.

n개를 원하는 m명이 k개를 원하도록 고치려면

(n-k) * m 개가 필요하다.

 

내가 갖고있는 사탕이 10개이므로

충분히 줄 수 있다.

10개에서 4개를 빼면 6개가 남는다.

 

그리고 이렇게 주고나면 저장했던 데이터는 다음과 같이 변한다.

{2 : 4}

 

그 다음으로 작업을 하려고하니, 이제 그 다음 키가 없다!

그래서 한 가지 아이디어를 떠올렸다.

 

어차피 처음에 0개를 원하는 사람은 아무도 없다. 반드시 사탕을 원한다.

그리고 사탕을 원하는만큼 받은 사람은 0개를 원하게 된다.

그래서 초기 값으로 {0 : 0} 쌍을 추가한다.

그리고 사탕을 원하는 만큼 받은 인원은 이 쌍에 저장하면 된다.

 

그럼 저장했던 데이터를 다음과 같이 바꿀 수 있다.

{2 : 4, 0 : 0}

 

이제 다시 아까의 작업을 반복해보자.

2개를 원하는 사람이 4명있다. 0개를 원하도록 하려면 몇개가 필요할까?

(2-0) * 4 = 8 개가 필요하다.

그런데 6개밖에 없다.

 

이때는 어떻게 해야할까?

균등하게 줄 수 있는 만큼 줘야한다.

초등학교 때 처음 나눗셈을 배웠던 기억을 떠올리면 쉽다.

 

균등하게 줄 수 있는 만큼 = 몫

6개를 4명이 균등하게 나눌 수 있을 몫을 구하면?

6 // 4 = 1 (파이썬 기준 몫을 구하는 연산자이다)

 

따라서 모두 균등하게 1씩 깐다.

그리고 사탕도 (몫 * 인원수) 만큼 깐다.

6 - 4 = 2

바뀐 데이터는 다음과 같다.

{ 1 : 4, 0 : 0 }

 

 이제 1개씩 원하는 4명에게 2개를 1개씩 배분해주면 끝이다.

그러면 변화한 자료구조는

{ 1 : 2, 0 : 2 }

 

구하는 값은

sum(키*키*인원수)

 

따라서 1*1*2 + 0*0*2 = 2 이다.

 

지금까지 말로 풀어쓴 것을 코드로 구현한 것이 다음과 같다.

# 2878 캔디캔디
import sys
candy, friend = map(int, input().split())
want = dict()
want[0] = 0
for _ in range(friend): # O(n)
    w = int(sys.stdin.readline())
    try:
        want[w] += 1
    except:
        want[w] = 1

want_key = sorted(list(want.keys()), reverse=True) # O(n log n) (최악)
for i in range(len(want_key)-1):
    need = (want_key[i] - want_key[i+1])*want[want_key[i]]  # 그 다음 캔디에 맞추기 위해 필요한 사탕 개수
    if candy >= need:
        candy -= need
        want[want_key[i+1]] += want.pop(want_key[i])
    else:
        _key = want_key[i] - candy//want[want_key[i]]
        pop_v = want.pop(want_key[i])
        try:
            want[_key] += pop_v
        except:
            want[_key] = pop_v

        want[_key] -= candy%pop_v
        try:
            want[_key-1] += candy%pop_v
        except:
            want[_key-1] = candy%pop_v
        break
answer = 0
for key in want.keys():
    answer += key*key*want[key]
    answer %= 18446744073709551616
print(answer)

 

try / except 구문을 활용하여 존재하지 않는 키에 접근할 경우,

새로 키를 만들도록 처리했다.

in 연산자로 키가 존재하는지 매번 체크하는 것보다 효율적이다.