[Python - 프로그래머스 Level2] 배상 비용 최소화

 

문제를 잘못이해하여 처음에는 순열과 조합을 이용해 계산하였다 하지만 너무 복잡한 식이 되어 이게 맞나? 헷갈려 다른분들의 블로그를 참고하여 다시 문제를 이해하고 코드를 짤 수 있었다. 매우 간단하게 for문을 이용하여 매번 for문 안에서 works 리스트를 정렬하며 맨 마지막 (즉 최대값을 갖는 item)을 1씩 빼주는 방식으로 구현할 수 있었다.

 

for문과 sort 활용 풀이 과정

def solution(no, works):
    for _ in range(no):
        works.sort()
        if works[-1] == 0:
            break
        works[-1] -= 1

    return sum([i*i for i in works])

 

정확도는 다 맞을 수 있었지만 효율성은 모두 시간 초과 나는 풀이법이였다. 그래서 다른 풀이과정을 보던 도중 아래 블로그를 참고하여 Heap 자료구조를 활용하여 다시 풀어보게 되었다.

파이썬의 Heapq 모듈은 MinHeap 기반으로 이를 가지고 MaxHeap으로 원소를 (-value , value) 원소로 지정하여 구현하는 방법이다.

works 리스트를 MaxHeap으로 변환 후, no 만큼 순회하며 최대값 -1 연산을 진행한 다음 1번째 원소를 가지고 제곱값을 모두 더하여 반환하였다. 

참고

[Algorithm] [Python] Programmers - 배상 비용 최소화

 

Heap 풀이 과정

def solution(no, works):
    MaxHeap = []
    if no >= sum(works):
        return 0
    for i in works:
        heapq.heappush(MaxHeap,(-i,i))
    for _ in range(no):
        i = heapq.heappop(MaxHeap)[1]-1
        heapq.heappush(MaxHeap,(-i,i))
    return sum(i[1]**2 for i in MaxHeap)

 

 

자료구조 Heap에 대해 부족한 점을 많이 느끼게 된 문제였다.. Heap 말고도 자료구조를 기초부터 탄탄히 다시 공부해야겠다고 생각하게 되었다.