SWEA 1208번 D3 [S/W 문제해결 기본] 1일차 - Flatten 파이썬

리뷰

또 출력할때 테스트 케이스 앞에 #을 넣어주지 않아 오답이 노출되었다... 후

문제 풀이

1. 상자 정보를 리스트로 받아오고 상자 정보에 부호가 -인 상자 정보 리스트를 추가로 생성한다.
2. 상자 정보를 최소힙으로 -인 상자 정보를 최대힙으로 만들어준다.
3. d가 0이 될때까지 while 루프를 돌아주며 만약 최대값과 최소값이 1이내일 경우 루프를 빠져나온다.
4. 테스트 케이스와 최대힙과 최소힙의 차이를 출력해 준다.

참고 사항

힙을 사용하지 않을 경우 시간 초과가 노출되는 것 같다.

정답 코드

def q1208():
    # SWEA 1208번 D3 [S/W 문제해결 기본] 1일차 - Flatten 파이썬
    import heapq

    for i in range(1, 11):
        d = int(input())
        lst1 = list(map(int, input().split()))
        lst2 = [-j for j in lst1]

        heapq.heapify(lst1)
        heapq.heapify(lst2)

        while d:
            min_val = heapq.heappop(lst1)
            max_val = -heapq.heappop(lst2)
            if max_val - min_val <= 1:
                print(max_val - min_val)
                break
            heapq.heappush(lst1, min_val + 1)
            heapq.heappush(lst2, -(max_val - 1))
            d -= 1
        print(f'#{i} {-heapq.heappop(lst2) - heapq.heappop(lst1)}')
q1208()