[LeetCode | Python] 1046. Last Stone Weight

문제 요약

스톤의 무게가 담긴 stones[i] 배열이 있을 때

1. 가장 무거운 두 개의 스톤을 선택,

2. 같으면 둘다 제거, 다르면 큰 돌이 (큰 값 - 작은 값)으로 변함

3. 최대 한 개의 돌이 남아있을 때까지 반복

4. 남은 돌의 무게 반환 (없으면 0)

 

최적 풀이 : max heap 이용

max heap을 이용하면 항상 가장 큰 무게의 두 돌이 항상 배열의 맨 앞에 있는 구조를 유지할 수 있다.

class Solution:
    def lastStoneWeight(self, stones: List[int]) -> int:
        stones = [-s for s in stones]
        heapq.heapify(stones)
        
        while len(stones) > 1:
            x = -heapq.heappop(stones)
            y = -heapq.heappop(stones)
            if x > y:
                heapq.heappush(stones, -(x-y))
        
        return -stones[0] if stones else 0

heapq는 기본적으로 min heap을 만들기 때문에, max heap을 만들기 위해서는

배열의 모든 요소를 음수로 바꾼 후 heapify 해준다.

 

참고로 "다르면 큰 돌이 (큰 값 - 작은 값)으로 변함" 이 문제 조건을 잘 보면,

heap의 관점에서 봤을 때 큰 값은 어차피 사라지는 것이기 때문에

heappop을 두번 해준 후 heappush로 (큰 값 - 작은 값)을 해주면 된다.

일반적인 arr처럼 큰 값에 값을 직접 대입하는 식으로 다루지 않도록 주의하자.

 

heapq 메서드는 힙 생성에 O(n), 삽입 삭제에 O(log n)이므로

위 풀이의 시간 복잡도는 O(n log n)이다. 

 

sorting으로 푼다면? - 시간복잡도

이 문제 역시 sorting으로 풀이가 가능하지만, 

heapq를 사용하면 시간 복잡도가 크게 줄어들기 때문에

heap을 사용할 수 있는 경우엔 heap을 사용하는 것이 좋다. 

class Solution:
    def lastStoneWeight(self, stones: List[int]) -> int:
        
        while len(stones) > 1:
            stones.sort()
            cur = stones.pop() - stones.pop()
            if cur:
                stones.append(cur)
                
        return stones[0] if stones else 0

 

리스트의 sort()는 Timsort(파이선 기본 정렬 알고리즘)을 사용하며 O(N log N)의 시간복잡도를 가진다.

while문 안에서 매번 sort()를 해주기 때문에

위 sorting 풀이의 시간 복잡도는 O(n^2 log n)이다.