[LeetCode] 2593. Find Score of an Array After Marking All Elements (Python)

 

난이도: Medium

 

문제 설명

---

You are given an array nums consisting of positive integers.

Starting with score = 0, apply the following algorithm:

• Choose the smallest integer of the array that is not marked. If there is a tie, choose the one with the smallest index.
• Add the value of the chosen integer to score.
• Mark the chosen element and its two adjacent elements if they exist.
• Repeat until all the array elements are marked.

Return the score you get after applying the above algorithm.

 

문제 예제

---

Example 1:

Input: nums = [2,1,3,4,5,2]
Output: 7
Explanation: We mark the elements as follows:
- 1 is the smallest unmarked element, so we mark it and its two adjacent elements: [2,1,3,4,5,2].
- 2 is the smallest unmarked element, so we mark it and its left adjacent element: [2,1,3,4,5,2].
- 4 is the only remaining unmarked element, so we mark it: [2,1,3,4,5,2].
Our score is 1 + 2 + 4 = 7.

Example 2:

Input: nums = [2,3,5,1,3,2]
Output: 5
Explanation: We mark the elements as follows:
- 1 is the smallest unmarked element, so we mark it and its two adjacent elements: [2,3,5,1,3,2].
- 2 is the smallest unmarked element, since there are two of them, we choose the left-most one, so we mark the one at index 0 and its right adjacent element: [2,3,5,1,3,2].
- 2 is the only remaining unmarked element, so we mark it: [2,3,5,1,3,2].
Our score is 1 + 2 + 2 = 5.

 

제한 사항

---

• 1 <= nums.length <= 10^5
• 1 <= nums[i] <= 10^6

 

✏️ Solution(솔루션)

---

class Solution(object):
    def findScore(self, nums):
        n = len(nums)
        check = [False for i in range(n)]
        nums_with_indexes = [(num, i) for i, num in enumerate(nums)]
        heapq.heapify(nums_with_indexes)

        answer = 0
        check_cnt = 0
        while check_cnt != n:
            min_num, idx = heapq.heappop(nums_with_indexes)
            if not check[idx]:
                answer += min_num
                check[idx] = True
                check_cnt+=1
                if idx-1 >= 0 and not check[idx-1]:
                    check[idx-1] = True
                    check_cnt+=1
                if idx+1 < n and not check[idx+1]:
                    check[idx+1] = True
                    check_cnt+=1

        return answer

 

 

나는 heapq를 사용해서 문제를 해결했다.

 

일단 nums의 요소와 idx를 튜플로 묶어 리스트에 저장하고 (ex. nums = [1,2,3] ==> nums_with_indexes = [(1,0), (2,1), (3,2)]) 이를 heapq로 만들었다.

 

이렇게 하면 매번 작은 값을 pop할 수 있고, 그 값의 인덱스도 알 수 있다.

 

그 후, mark를 하기 위해 check라는 리스트를 만들었고, 인덱스가 mark가 된건지 아닌지 확인했다.

 

예를 들어 heappop해서 나온 값이 min_num = 1, idx=3 이라고 한다면 일단 check[idx]가 True인지 확인하고 아니라면 answer에 min_num을 더해주고 check[idx]와 check[idx-1], check[idx+1]을 True로 변환해 주었다.

 

그리고, check에 모든 value가 True라고 한다면 연산을 마치고 answer를 return하여 정답을 맞출 수 있었다.

 

처음에는 while False not in check: 로 했는데 Time Limit Exceed 결과가 나와서 check_cnt라는 변수를 만들었고, True가 생길때마다 +1 해주었다.

 

이제 이 check_cnt 변수가 nums의 길이와 같아지면 반복문을 종료하게 수정했고, Time Limit Exceed 없이 정답을 맞출 수 있었다.

 

다 풀고 다른 사람들의 풀이를 보니, 굳이 heapq를 사용안하고 정렬만해도 쉽게 풀 수 있는 문제인 것 같다..

---

2593. Find Score of an Array After Marking All Elements