[Python] 백준/BOJ 23757번: 아이들과 선물 상자 (Silver 2)

💻 Problem

문제 보러 가기

 

상훈이는 N개의 선물 상자를 가지고 있다. 선물 상자에는 현재 담겨있는 선물의 개수가 적혀있다.

선물을 받을 아이들이 M명 있다. 아이들은 각자 1에서 M까지의 서로 다른 번호를 하나씩 부여받았다.

 1번 아이부터 M번 아이까지 한 번에 한 명씩, 현재 선물이 가장 많이 담겨있는 상자에서 각자 원하는 만큼 선물을 가져간다. 이때, 앞서 누군가 선물을 가져갔던 선물 상자에서 또다시 가져가도 상관없다.

하지만 상자에 자신이 원하는 것보다 적은 개수의 선물이 들어있다면, 선물을 가져가지 못해 실망한다.

상훈이는 한 명이라도 실망하지 않고 모두가 선물을 가져갈 수 있는지 궁금하다.

 

💡 Approach

heapq 모듈을 활용하여 최대 힙을 구현했다.

최대 힙의 경우 최소 힙과 동일하게 구현하고 원소를 넣을 때 음수로 넣으면 된다.

음수로 넣었으니 원소를 빼서 사용할 때는 다시 양수로 만들어줘야 한다.

 

입력을 처리한 후에는 아래와 같은 로직으로 구현했다.

1. 아이들을 한 명씩 처리한다
2. 현재 선물이 가장 많이 담겨있는 상자를 꺼낸다
3. 상자에 자신이 원하는 것보다 크거나 같은 개수의 선물이 들어있다면 선물을 가져간다
3-1. 이때, 아이가 선물을 가져가고도 상자에 선물이 남아있으면 다시 present에 push 한다
4. 상자에 자신이 원하는 것보다 적은 개수의 선물이 들어있다면, 선물을 가져가지 못한다 (즉시 종료)

 

✏️ Solution

import sys
import heapq
input = sys.stdin.readline

n, m = map(int, input().split()) # 선물 상자의 수, 아이들의 수
present = [-int(num) for num in input().split()] # 선물 상자에 들어있는 선물의 개수
children = list(map(int, input().split())) # 각 아이가 원하는 선물의 개수
heapq.heapify(present)

for child in children:
    x = -heapq.heappop(present)
    # 상자에 자신이 원하는 것보다 크거나 같은 개수의 선물이 들어있다면 선물을 가져간다
    if x >= child:
        if (x - child) > 0:
            heapq.heappush(present, child - x)
    # 상자에 자신이 원하는 것보다 적은 개수의 선물이 들어있다면, 선물을 가져가지 못한다
    else:
        print(0)
        exit()
print(1)

3-1(x - child > 0) 코드는 없어도 가능하다.

선물이 남아 있을 때만 넣어야 된다고 생각했는데 최대 힙이여서 0을 넣어도 계산할 일이 없어 상관없다.

 

시간복잡도: O(n + m log n)

heapify의 시간복잡도가 O(n)이고, 아이들의 선물을 처리하는 반복문 코드가 O(m log n)이다.