[Python 문제 풀이] 프로그래머스 '광물 캐기'

Programmers Lv.2 파이썬 '광물 캐기' 풀이 및 코드

마인크래프트..?
게임 설정 상 광물로 곡괭이를 만들 수 있으며, 더 좋은 재료로 만들어진 곡괭이일 수록 광물을 더 쉽게 채굴할 수 있다.

 

조건파악

이 문제는 조건이 매우 중요하다. 
1. 광물은 무조건 5개씩 연속으로 캔다.
2. 5개를 채굴할 때마다 곡괭이를 고른다.
 
어떤 곡괭이를 먼저 쓸 것인가?
순서대로 채굴해야 하기 때문에 minerals를 재정렬하는 것은 불가능하다. 하지만 5개 채굴시마다 곡괭이를 바꾸기 때문에 크기가 5인 청크로 묶는다면 서로 다른 청크는 서로 독립적이기 때문에 정렬을 수행해도 결과에 영향을 주지 않는다!
 

picks는 원소가 3개 뿐이며, dia,iron,stone은 각각 해당 재료로 만들어진 곡괭이의 수를 의미한다.
 
곡괭이가 어느 청크에 사용될지 정하기 위해 minerals의 모든 원소를 확인할 필요가 있는데 다행히 최대 길이는 50이다. 시간복잡도는 크게 신경 안 쓰고 풀어도 되겠다.
 

접근

다시 한 번 정리하자면, minerals를 크기가 5인 청크로 묶어서 생각할 수 있다. 어떤 청크를 어떤 곡괭이로 캘 것인가가 문제를 푸는 방법이다.
 
피로도를 최소화하기 위해서는 청크 안에 귀한 광물이 얼마나 포함되어있는지 파악할 필요가 있다. 그러고 나서 귀한 광물이 포함된 청크들을 먼저 오게 정렬하고 귀한 곡괭이를 먼저 쓰면 피로도가 최소가 된다.
 
다행히 각 광물을 채굴하는데 드는 피로도는 5의 제곱으로 증가하기 때문에 diamond가 하나라도 있는 것이 iron이 5개 있는 청크보다 더 높은 피로도를 요구한다. 따라서 각 광물의 종류를 독립적으로 보고 정렬해도 무방하다. 

아주 단순한 정렬 문제이다.
 

구현

정렬이 있어서 시간복잡도가 O(nlogn)이지만 크기 5인 청크로 묶었기 때문에 n = n/5이다. 생각보다 실행시간을 짧게 구현할 수 있다.

def solution(picks, minerals):
    # 광물의 수가 (곡괭이의 수) x 5 보다 많다면,
    # 채굴 가능한 총 광물의 개수를 자원 (곡괭이의 수) x 5로 제한한다.
    if sum(picks) * 5 < len(minerals):
        minerals = minerals[:sum(picks) * 5]

    # 광물을 크기가 5인 청크로 분할하고 각 청크에 포함된 종류별 광물의 개수를 센다.
    # 광물의 개수를 기준으로 내림차순으로 정렬한다.
    counted = scan_minerals(minerals)

    # 정렬 방법에 따라 곡괭이의 개수를 줄여가며 피로도를 계산한다.
    answer = calculate_fatigue(counted, picks)
    return answer

def scan_minerals(minerals):
    i = 0
    counted = []
    flag = True
    while flag:
        target = []
        if i + 5 < len(minerals):
            target = minerals[i:i + 5]
        else:
            target = minerals[i:]
            flag = False
        dias, irons, stones = target.count('diamond'), target.count('iron'), target.count('stone')
        counted.append([dias, irons, stones])
        i += 5
    counted.sort(key=lambda _: (-_[0], -_[1]))
    return counted

def calculate_fatigue(counted, picks):
    result = 0
    for target in counted:
        if picks[0] > 0:
            picks[0] -= 1
            result += sum(target)
        elif picks[1] > 0:
            picks[1] -= 1
            result += target[0] * 5 + target[1] + target[2]
        elif picks[2] > 0:
            picks[2] -= 1
            result += target[0] * 25 + target[1] * 5 + target[2]
        else:
            break
    return result