백준 23757번 파이썬

문제 링크

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

23757번: 아이들과 선물 상자

 

시간 제한 / 메모리 제한

1 초

1024 MB

 

문제

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

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

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

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

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

입력

첫째 줄에 선물 상자의 수 N과 아이들의 수 M이 공백을 사이에 두고 주어진다. (1 ≤ M ≤ N ≤ 10^5)

둘째 줄에 각 선물 상자에 들어있는 선물의 개수 c1,c2,…,cN 이 공백을 사이에 두고 주어진다. (1 ≤ ci ≤10^5)

셋째 줄에 아이들의 번호 순으로 각 아이가 원하는 선물의 개수 w1,w2,…,wM이 공백을 사이에 두고 주어진다. (1 ≤ wi ≤ 10^5)

출력

모든 아이들이 실망하지 않고 각자 원하는 만큼 선물을 가져갈 수 있으면 1을, 그렇지 않으면 0을 출력한다.

예제 입력

4 4
4 3 2 1
3 1 2 1

예제 출력

1

 

풀이 코드

import heapq as h

n, m = map(int, input().split())
box = list(map(int, input().split()))
child = list(map(int, input().split()))
boxH = []

for b in box:
    h.heappush(boxH, -b)

for c in child:
    b = -h.heappop(boxH)
    if b > c:
        h.heappush(boxH, -(b - c))
    elif b == c:
        continue
    else:
        print(0)
        exit()
print(1)

 

문제에 "현재 선물이 가장 많이 담겨있는 상자에서"라는 조건이 주어졌기에 최대 힙 자료구조를 사용할 수 있다. python heapq는 기본적으로 min heap(최소 힙)이기 때문에 최대 힙을 이용하려면 heappush와 heappop 연산 시 부호를 바꿔주어야 한다.

 

1번 아이부터 m번 아이까지 순서대로 가장 큰 선물 상자에서 선물을 가져가기 위해 각 아이의 순서가 오면 우선 최대 힙에서 가장 큰 값을 pop한다. 각 아이가 받고 싶어하는 선물의 수 만큼 선물을 가져가고도 선물 상자에 선물이 남았다면 이를 최대 힙에 다시 추가해주고, 정확히 선물의 갯수와 아이가 받고 싶어하는 선물의 수가 일치한다면 그 상자는 소모된다(최대 힙에 추가하지 않는다). 만약 최대 힙을 pop한 결과가 현재 순서 아이가 받고 싶어하는 선물의 수보다 작다면, 최대 힙에 이보다 큰 값이 없고, 이 아이는 선물을 받을 수 없다는 뜻이기에 0을 출력하고 프로그램을 종료한다. 그리고 모든 아이가 선물을 받을 수 있었다면 1을 출력한다.