문제 링크

문제 간단 해석

가장 작은 값들을 골라 특정 계산식으로 섞는 문제.
계속해서 작은 값들을 골라내야 하므로, heap을 사용해야 한다.

Java

풀이 코드

import java.util.Queue;
import java.util.PriorityQueue;
import java.util.Arrays;
import java.util.List;
import java.util.LinkedList;
import java.util.stream.Collectors;

class Solution {
    public int solution(int[] scoville, int K) {
        List<Integer> list = Arrays.stream(scoville).boxed().collect(Collectors.toList());
        Queue<Integer> queue = new PriorityQueue<>(new LinkedList(list));
        int A, B, count = 0;
        
        while (true) {
            try {
                A = queue.remove();
                
                if (A >= K)
                    return count;
                
                B = queue.remove();
                queue.add(A + B  * 2);
                count++;
                
            } catch (Exception e) {
                return -1;
            }
        }
    }
}

해석

List<Integer> list = Arrays.stream(scoville).boxed().collect(Collectors.toList());
Queue<Integer> queue = new PriorityQueue<>(new LinkedList(list));

int Array인 scoville을 Integer List로 만들었다.
그 후 우선순위 큐를 선언하고 방금 만든 리스트를 LinkedList로 만들어 넣어주었다.

while (true) {
    try {
        A = queue.remove();

        if (A >= K)
            return count;

        B = queue.remove();
        queue.add(A + B  * 2);
        count++;

    } catch (Exception e) {
        return -1;
    }
}

while 안에서 가장 작은 값 2개를 뽑는다.
만약 처음에 뽑은 값이 K보다 크거나 같을 경우, 조건 충족이기에 연산 횟수인 count를 반환한다.
두번째 값을 뽑고, 문제에서 주어진 조건대로 잘 섞어 다시 우선순위 큐에 넣어준다.
그 후 count를 증가시킨다.
(해당 문제에서는 catch에서 return하기에 문제되지 않지만, 그래도 try catch문 이후로 count를 옮겨야 더 이쁜 코드가 나올 것 같다.)

만약 값을 뽑을 수 없는 경우가 생겼을 때(NoSuchElementException 발생), 주어진 조건으로는 K를 충족시킬 수 없는 경우이므로 -1을 반환한다.

Python

풀이 코드

import heapq

def solution(scoville, K):
    heapq.heapify(scoville)
    mix = 0

    while len(scoville) > 1:
        scov_min = heapq.heappop(scoville)

        if scov_min >= K:
            return mix

        scov_sec_min = heapq.heappop(scoville)
        heapq.heappush(scoville, scov_min + scov_sec_min * 2)
        mix += 1

    return mix if scoville[0] >= K else -1

해석

heapq.heapify(scoville)

scoville을 우선순위 큐 형태로 내부 정렬.

while len(scoville) > 1:
    scov_min = heapq.heappop(scoville)
    
    if scov_min >= K:
        return mix
    
    scov_sec_min = heapq.heappop(scoville)
    heapq.heappush(scoville, scov_min + scov_sec_min * 2)
    mix += 1

가장 작은 값 하나를 빼오고, 그 값이 K 이상이라면 섞은 횟수를 반환한다.
두번째로 작은 값을 또 빼오고, 섞은 값을 다시 넣어준다.
그 후 섞은 횟수를 증가시킨다.

return mix if scoville[0] >= K else -1

만약 값이 한개밖에 남지 않았을 경우, 해당 값이 K 이상이면 섞은 횟수, 아니라면 -1을 반환한다.