heapq 모듈을 사용한 풀이입니다.

heap을 이용해 후보들 중 가장 적합한 후보를 추출할 수 있습니다.
자세한 사용법은 포스트 헤더의 learn more을 클릭하시면 볼 수 있습니다.

 

문제 설명

XX 게임의 유저들이 보스 몬스터를 사냥하려고 팀을 만들었습니다. 그리고 팀에 속한 캐릭터에 아이템을 사용해 공격력을 높이려 합니다.

이 게임의 아이템은 캐릭터의 공격력은 높이고 체력을 낮춥니다. 그래서 아이템을 적절히 사용해 팀의 공격력을 최대한 끌어올리려 합니다. 캐릭터별로 아이템을 사용할지 말지는 자유지만, 아이템을 사용한 캐릭터는 체력이 반드시 100 이상 남아야 합니다. 또, 한 캐릭터에 아이템 하나씩만 사용할 수 있으며, 사용한 아이템은 사라집니다.

예를 들어 캐릭터의 체력이 [200, 120, 150]이고 아이템의 효과는 다음과 같습니다.

|높여줄 공격치 |낮추는 체력| |—|—| |30 |100| |500| 30| |100| 400|  

이때 팀의 공격력을 최대로 올리려면 첫 번째 캐릭터에 첫 번째 아이템을, 세 번째 캐릭터에 두 번째 아이템을 사용하면 됩니다.

캐릭터들의 체력을 담은 배열 healths와 아이템별 효과를 담은 이차원 배열 items가 solution 함수의 매개변수로 주어질 때, 팀의 공격력을 최고로 끌어올리려면 어떤 아이템을 사용해야 하는지 return 하도록 solution 함수를 완성해주세요.

 

제한 조건

healths의 길이는 1 이상 10,000 이하입니다. healths의 원소(캐릭터의 체력)는 1 이상 1,000,000 이하인 자연수입니다. items의 길이는 1 이상 5,000 이하입니다. items에는 아이템이 [올려줄 공격력, 낮출 체력]이 번호 순서대로 들어있습니다. 아이템 번호는 1번 부터 시작합니다. items[i]에는 i + 1번 아이템이 [올려줄 공격력, 낮출 체력]이 들어있습니다. 아이템이 올리는 공격력은 1 이상 500,000 이하인 자연수입니다. 아이템이 내리는 체력은 1 이상 500,000 이하인 자연수입니다. 아이템 번호는 오름차순으로 정렬해 return 해주세요. 올려주는 공격력이 같은 아이템은 없습니다. 아이템을 사용하는 방법이 여러 가지라면, 그러한 방법 중 아무거나 하나를 return 해주세요. 단, 아이템 번호는 오름차순으로 정렬되어 있어야 합니다.

입출력 예 |healths |items |return| |—|—|—| |[200,120,150] |[[30,100],[500,30],[100,400]]| [1,2]| |[300,200,500] |[[1000, 600], [400, 500], [300, 100]] |[3]|

 

 

 

문제 풀이

import heapq
from collections import deque

def solution(healths, items):
    healths.sort()
    items = [(*item,index) for index, item in enumerate(items, 1)]
    ATTACK, HP, IDX = 0, 1, 2

    items = deque(sorted(items, key=lambda x:x[HP]))
    
    answer = []
    cand = []
    for health in healths:
        # 유저에 대한 item 후보를 heap에 쌓아 줍니다. 
        # 매 유저에 대해 maxheap으로 heap이 갱신됩니다. 
        while items and health - items[0][HP] >= 100:
            item = items.popleft()
            heapq.heappush(cand, (-item[ATTACK], item[IDX]))
        # 힙에 아이템이 있는 경우 가장 공격력이 높은 아이템 하나를 후보에 넣습니다. 
        if cand:
            answer.append(heapq.heappop(cand)[1])
    return sorted(answer)

이처럼 heap을 이용하면 후보들을 계속 sorting 해가면서 우선순위 큐를 활용할 수 있습니다.

 

해설 답안

  1. 체력을 오름차순으로 정렬하고 O(h log h)
  2. 루프를 돌면서 HP 100 이상이라는 조건을 만족하는 것들 중
    가장 공격력을 많이 올려주는 것을 선택 O(i)
  3. 선택하고나면 해당 아이템 제거 O(i)

정렬은 제외하더라도 루프와 아이템 제거에서 O($i^2$) 만큼 시간복잡도가 소요되기 때문에 효율성 테스트에서 틀립니다.
그러므로 힙을 사용해 해결을 진행합니다.

from heapq import heapify, heappush, heappop

def solution(healths, items):
    healths.sort() # 체력을 오름차순으로 정렬
    items = sorted([(item[1], item[0], index+1) for index, item in enumerate(items)]) # 깎는 체력 순으로 정렬
    answer = []
    heap = []

    for health in healths: # 제일 작은 체력부터 루프
        while items: # 아이템 루프
            debuff, buff, index = items[0] # 가장 깎는 체력이 낮은 아이템
            if health - debuff < 100: # 조건이 안맞으면 루프 종료
                break
            items.pop(0) # 어차피 다음 체력에선 무조건 가능하기 때문에 이 시점에서 제거합니다. 여기서 deque를 사용해도 되겠죠?
            heappush(heap, (-buff, index)) # 힙에 추가합니다
        if heap: # 힙이 비어있지 않으면
            _, index = heappop(heap)
            answer.append(index) # 답 리스트에 추가합니다.

    return sorted(answer)

 

 

참고문헌

https://github.com/tjdud0123/daily_algorithm https://www.daleseo.com/python-heapq/

댓글남기기