[Programmers] 더 맵게 (Python, Java)

더 맵게

programmers.co.kr/learn/courses/30/lessons/42626

 

숫자 목록이 주어지고 거기서 가장 낮은 숫자와, 그 다음 숫자 * 2를 더한 값을 다시 목록에 넣을 수 있는 횟수를 리턴하는 문제이다.

 

목록을 값에 추가하고, 나올 때 무조건 작은 값이 먼저 나와야 된다는 조건이 있기 때문에 min-heap 을 생각하면 그 다음엔 알맞은 자료구조를 선택하여 구현하면 된다.

 

파이썬에선 heap, 자바에선 PriorityQueue 를 사용할 수 있다.

(만약 큰 수대로 나와야 한다면 값을 추가할 때 -1 을 곱해서 사용하면 된다.)

heap 사용하여 풀이 (Python)

from heapq import heappush, heappop


def solution(scoville, K):
    answer = 0
    h = []
    for s in scoville:
    	# 힙에 값 추가
        heappush(h, s)
    
    # 힙에 있는 원소 개수
    # len(h) 는 오래 걸리는 연산이 아니지만 +,- 할 수 있는 값이여서 변수 remain 을 따로 둠
    remain = len(h)
    while remain > 1:
        s1 = heappop(h)
        # 힙에서 원소 개수 뺏기 때문에 -1
        remain -= 1
        # 힙에서 나온 값은 가장 작은 값이다.
        # 이 값이 K 보다 작은 경우
        if s1 < K:
        	# 두번 째로 스코빌 지수가 낮은 음식을 꺼낸다.
            # 음식을 꺼내고, 다시 집어 넣을 것이기 때문에 remain 은 값 변화가 없어서 무시함
            s2 = heappop(h)
            # 음식을 섞은 횟수 (정답) 증가
            answer += 1
            # 다시 음식 넣기
            heappush(h, s1 + (s2 * 2))
        else:
        	# 힙에서 나온 값이 K 이상이다.
            # 이러면 그 뒤의 값들은 무조건 K 이상이므로
            # 모든 원소가 K 보다 크다고 볼 수 있다. (음식의 스코빌을 섞을 필요가 없음)
            # 그래서 바로 정답을 리턴한다.
            return answer
    
    # remain 이 1일 때, 즉 힙에 남아 있는 음식이 한개일 때 여기로 오게 됨
    # 이 때 힙에 있는 유일한 음식이 K 이상의 스코빌 지수를 갖는다면
    if h and h[0] >= K:
    	# 정답을 리턴
        return answer
    # 힙에 남아 있는 음식이 K 미만인 경우엔
    # 더 이상 섞을게 없다는 의미이므로 -1 리턴
    return -1

파이썬에서 heap 은 기본이 min-heap 이다.

그래서 처음에 나온 값이 가장 낮은 값 인것이 보장이 된다.

 

그래서 원소 두개를 빼고, 문제에서 주어진 식으로 (a[i] + (a[i+1] * 2)) 계산하여 다시 힙에 넣는다.

 

정답을 만드는 조건은, 다시 힙에 push 했을 때, 즉 두 음식의 스코빌 지수를 섞을 때이다.

 

이렇게 계속 음식을 섞다가 (힙에 넣고 빼고 하다가), 더 이상 섞을 게 없다는 의미는

 

- 힙에 원소가 하나만 남거나 (1)

- 힙에서 뺀 값이 스코빌 K 이상이 되는 경우다. (2)

 

(2) 는 힙에서 뺀 값이 K 이상이면 바로 answer 를 리턴 하면 된다.

 

(1) 은 힙에 남아 있는 개수를 확인하며 while 루프를 돈다. while 루프를 다 도는 경우는 힙에 원소가 하나만 남는 경우로 이 때 하나만 남은 값이 K 미만이라면 음식을 모두 섞었음에도 모든 원소가 K 이상이 될 수 없다는 의미이다. 이 조건을 확인하여 -1 을 리턴하면 된다.

 

heap 사용하여 풀이 (Java)

자바로도 비슷하게 풀 수 있다. PriorityQueue 자료구조를 사용하면 된다.

import java.util.PriorityQueue;


class Solution {
    public int solution(int[] scoville, int K) {
        int answer = 0;
        PriorityQueue<Integer> heap = new PriorityQueue<>();

        for (int s: scoville) {
            heap.add(s);
        }

        while (heap.size() > 1) {
            int s1 = heap.poll();
            if (s1 >= K) {
                return answer;
            }
            int s2 = heap.poll();
            answer++;
            heap.add(s1 + (s2 * 2));
        }
        
        if (!heap.isEmpty() && heap.peek() > K) {
            return answer;
        }

        return -1;
    }
}