Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

파이썬으로 해결하는 마지막 남은 돌의 무게(Last Stone Weight) 문제

문제 소개

양의 정수 무게를 가진 돌 여러 개가 주어졌다고 가정해 보겠습니다. 매 턴마다 가장 무거운 두 개의 돌을 골라 서로 부딪쳐 부수는 작업을 반복합니다. 두 돌의 무게를 각각 x와 y(x ≤ y)라고 할 때, 충돌 결과는 다음 두 가지 중 하나입니다.

  • x = y인 경우: 두 돌 모두 완전히 파괴됩니다.
  • x ≠ y인 경우: 무게가 x인 돌은 완전히 파괴되고, 무게가 y인 돌은 새로운 무게 y − x를 갖게 됩니다.

이 과정이 끝나면 최대 1개의 돌만 남습니다. 우리가 구해야 하는 값은 바로 이 마지막 돌의 무게이며, 만약 돌이 하나도 남지 않았다면 0을 반환하면 됩니다.

예시로 이해하기

돌의 무게 배열이 [2, 7, 4, 1, 8, 1]일 때를 단계별로 살펴보겠습니다. 최종 결과는 1입니다.

  1. 가장 무거운 8과 7을 선택해 충돌 → 무게 1인 돌 생성, 배열은 [2, 4, 1, 1, 1]
  2. 4와 2를 선택해 충돌 → 무게 2인 돌 생성, 배열은 [2, 1, 1, 1]
  3. 2와 1을 선택해 충돌 → 무게 1인 돌 생성, 배열은 [1, 1, 1]
  4. 무게 1인 두 돌을 선택해 충돌 → 두 돌 모두 파괴, 배열은 [1]

따라서 마지막에 남는 돌의 무게는 1입니다.

해결 전략

이 문제는 단순한 시뮬레이션으로 해결할 수 있습니다. 핵심 아이디어는 다음과 같습니다.

  • 배열 W가 비어 있다면 0을 반환합니다.
  • W에 원소가 하나뿐이라면 그 값을 그대로 반환합니다.
  • W에 원소가 두 개 이상 남아 있는 동안 다음을 반복합니다.
    • W를 오름차순으로 정렬합니다.
    • s1 := 마지막 원소(가장 무거운 돌), s2 := 뒤에서 두 번째 원소(두 번째로 무거운 돌)
    • s1 = s2라면 두 원소를 모두 제거합니다.
    • 그렇지 않다면 s1 := |s1 − s2|로 갱신하고, 마지막 원소를 제거한 뒤 새로운 마지막 원소를 s1 값으로 덮어씁니다.
  • 반복이 끝난 후 W에 원소가 하나 남아 있다면 그 값을, 비어 있다면 0을 반환합니다.

파이썬 구현

다음 구현을 통해 더 자세히 이해해 보겠습니다.

class Solution(object):
    def lastStoneWeight(self, stones):
        """
        :type stones: List[int]
        :rtype: int
        """
        if len(stones) == 0:
            return 0
        if len(stones) == 1:
            return stones[0]
        while len(stones) > 1:
            stones.sort()
            s1, s2 = stones[-1], stones[-2]
            if s1 == s2:
                stones.pop(-1)
                stones.pop(-1)
            else:
                s1 = abs(s1 - s2)
                stones.pop(-1)
                stones[-1] = s1
        if len(stones):
            return stones[-1]
        return 0

ob1 = Solution()
print(ob1.lastStoneWeight([2, 7, 4, 1, 6, 1]))

실행 결과

입력:

[2, 7, 4, 1, 6, 1]

출력:

1

성능 개선: 최대 힙(Heap) 활용하기

위 구현은 매 반복마다 정렬을 수행하므로 시간 복잡도가 O(n² log n)입니다. 파이썬의 heapq 모듈은 기본적으로 최소 힙만 지원하지만, 무게에 음수를 취해 저장하면 최대 힙처럼 활용할 수 있습니다. 이 방식을 사용하면 각 연산을 O(log n)에 처리할 수 있어 전체 시간 복잡도를 O(n log n)까지 줄일 수 있습니다.

import heapq

class Solution(object):
    def lastStoneWeight(self, stones):
        heap = [-s for s in stones]
        heapq.heapify(heap)
        while len(heap) > 1:
            s1 = -heapq.heappop(heap)
            s2 = -heapq.heappop(heap)
            if s1 != s2:
                heapq.heappush(heap, -(s1 - s2))
        return -heap[0] if heap else 0

마무리

'마지막 남은 돌의 무게' 문제는 조건에 따라 두 돌을 충돌시키는 과정을 그대로 시뮬레이션하면 되는 직관적인 문제입니다. 정렬을 활용한 구현은 이해하기 쉽지만, 데이터 크기가 커질 경우 힙을 사용한 접근이 훨씬 효율적입니다. 두 가지 방법을 모두 익혀두면 코딩 테스트에서 유사한 유형의 문제를 빠르게 해결하는 데 큰 도움이 될 것입니다.