[리트코드/leetcode/python] 2593. Find Score of an Array After Marking All Elements

이미지를 누르면 문제 링크를 보실 수 있습니다

이미지를 누르면 문제 링크를 보실 수 있습니다

문제의 로직은 간단했습니다. 주어진 배열 nums에서 마크되지 않은 최소값을 더해가며 총 합을 출력하면 되는 문제였습니다. 여기서 마크되지 않았다는 것은 문제에서 주어진 예시를 보면 잘 이해가 될 것입니다.

 

쉬운 문제라고 생각했지만, 처음에는 풀지 못했습니다. 

제가 처음에 풀었던 코드는 아래와 같습니다.

class Solution:
    def findScore(self, nums: List[int]) -> int:

        res = 0
        mark = [0 for _ in range(len(nums))]

        tmp = nums[:]

        while len(tmp)!=1:
            
            res += min(tmp)
            target = tmp.index(min(tmp))
            print("index:",target,"min value:",min(tmp))
            print("res:",res)
            # mark[target] = 1
            if target == 0 and :
                del tmp[target+1]
                del tmp[target]
            elif target == (len(tmp)-1):
                del tmp[target-1]
                del tmp[target-1]
            else:
                del tmp[target+1]
                del tmp[target-1]
                del tmp[target-1]
            
            print("current_tmp:",tmp)
        
        return res+sum(tmp)

굉장히 뭔가 길지만, 해당 테스트케이스에서 오류를 마주하게 됐고 어디가 잘못됐는지 알 수 있었습니다.

[10,44,10,8,48,30,17,38,41,27,16,33,45,45,34,30,22,3,42,42]

 

제가 해결한 로직은 이러했습니다. 해당 문제에서 '마크'한다는 것은 주어진 배열에서 특정 값을 방문한다는 의미로 해석하면 됐습니다. 예를 들어, 주어진 배열에서 최소값을 찾은 후에, 그 값의 왼쪽과 오른쪽 값을 각각 마크하면, 마크되지 않은 원소들 중에서 최소값을 구해야 하는 원리가 반복되어 결국에 나중에는 마크되지 않은 원소는 하나만 남게 됩니다.

 

저는 결국 이 문제에서 사용되지 않은 원소이기에, 그냥 삭제하면 되는 줄 알았지만, 해당 테스트케이스에서 그 생각의 허점이 발견됩니다.

[10,44,10,8,48,30,17,38,41,27,16,33,45,45,34,30,22,3,42,42]

# min_val = 3 and mark it and left, right value
[10,44,10,8,48,30,17,38,41,27,16,33,45,45,34,30,(22),(3),(42),42]

# min_val = 8 and mark it and left, right value
[10,44,(10),(8),(48),30,17,38,41,27,16,33,45,45,34,30,(22),(3),(42),42]

# min_val = 10 and mark it and left, right value
# but, this time there is no left value(the out of the index)
[(10),(44),(10),(8),(48),30,17,38,41,27,16,33,45,45,34,30,(22),(3),(42),42]

# min_val = 16 and mark it and left, right value
[(10),(44),(10),(8),(48),30,17,38,41,(27),(16),(33),45,45,34,30,(22),(3),(42),42]

# min_val = 17 and mark it and left, right value
[(10),(44),(10),(8),(48),(30),(17),(38),41,(27),(16),(33),45,45,34,30,(22),(3),(42),42]

# min_val = 30 and mark it and left, right value
# and the error ocuurs here
[(10),(44),(10),(8),(48),(30),(17),(38),41,(27),(16),(33),45,45,34,30,(22),(3),(42),42]

처음 생각한 대로, 마크된 원소를 그냥 delete해버리면, 위 테스트케이스에서 min_val 가 30일 때, 오류가 발생합니다.

최소값이 발견된 위치 기준으로 왼쪽과 오른쪽값을 마크해야 하는데, 원소를 삭제해버리면 왼쪽 원소 '34'는 그렇다 쳐도,  '42'라는 원소가 '30' 오른쪽에 있는 것으로 인식 되고, 그대로 '42'가 삭제되는 상황이 발생하게 됩니다.

인덱스가 중요한 키포인트가 될 수 있음을 확인할 수 있습니다.

 

그렇게, 코드를 다시 짰지만, 시간초과가 나는 코드가 됐습니다. 

두 번째 시도로 만든 코드는 아래와 같습니다.

class Solution:
    def findScore(self, nums: List[int]) -> int:

        res = 0

        mark = [False for _ in range(len(nums))]

        while all(mark) != True:
            
            min_val = 100001
            for idx,num in enumerate(nums):
                if mark[idx] == False:
                    min_val = min(min_val,num)
                    if min_val == num:
                        target = idx

            res += min_val

            if target == 0 and (mark[target] == False):
                mark[target+1] = True
                mark[target] = True

            elif target == (len(tmp)-1) :
                mark[target-1] = True
                mark[target] = True

            elif 0 < target < (len(tmp)-1) :
                mark[target-1] = True
                mark[target] = True
                mark[target+1] = True
        
        return res

다시봐도 시간초과가 안 날 수 없는 코드입니다. 하지만, 이렇게 인덱스를 체크하면서 최소값을 찾아나가는 로직은 잘 정립된 과정이었습니다. 이제는 모든 위치가 다 mark가 되어 있는지와 최소값을 찾는 루트가 조금 더 빠른 로직을 짜야했습니다.

해답은 heapq 였습니다. heapq 라이브러리는 그 자체로 min_heap의 구조이기 때문에, 해당 라이브러리를 사용해야만 했습니다.

 

수정한 최종 코드는 아래와 같습니다.

import heapq

class Solution:
    def findScore(self, nums: List[int]) -> int:

        res = 0
        tmp = [(num,idx) for idx,num in enumerate(nums)]
        heapq.heapify(tmp)

        marked = set()
        
        while tmp:
            num, idx = heapq.heappop(tmp)
            if idx not in marked:
                res += num
                marked.add(idx)
                marked.add(idx-1)
                marked.add(idx+1)

        return res

tmp라는 리스트에 (num,idx)의 튜플 형태로 nums 리스트의 인덱스와 원소들을 저장합니다.

이 코드는 최종 코드까지 오면서 시도해봤던 코드이기도 합니다.

여기서 중요한 것은 이 tmp리스트를 heapify를 통해 힙큐화 시켰다는 것입니다. 이렇게 되면, 해당 tmp리스트는 min heap 구조를 띄게 되는데, 결국 이 tmp가 빌 때 까지 계속해서 최소값을 찾아가며 순회를 하면 됩니다.

여기서 중요한 두 번 째는 set을 사용했다는 것인데, 기존에 마크를 했는지 하지 않았는지 체크하기 위한 용도입니다.

똑같은 최소값을 찾아도, 해당 최소값이 이미 마크했던 최소값이라면, 똑같은 값을 가진 다른 위치의 값을 찾아갈 수 있겠죠.

 

로직을 해결하는 방법은 heapq와 set의 혼용이었습니다.

이 외에도 굳이 이 방법을 사용하지 않아도 되지만, 개인적으로 heapq를 다시한번 써볼 수 있고 배울 수 있는 문제였습니다.

 

이상으로 heapq과 관련한 2593. Find Score of an Array After Marking All Elements 문제를 정리해봤습니다.