https://school.programmers.co.kr/learn/courses/30/lessons/154539
문제요약
정수 배열에서 각 원소들에 대해 자신보다 뒤에 있고 자신보다 큰 수 중에, 가장 가까운 수를 뒷 큰수라고 정의한다. 뒷 큰수가 존재하지 않는 경우, -1로 정의한다. 정수로 이루어진 배열 numbers가 있을 때, 모든 원소의 뒷 큰수를 차례로 담은 배열을 return하는 함수를 만들어 보자.
제한사항
- 4 ≤ numbers의 길이 ≤ 1,000,000
- 1 ≤ numbers[i] ≤ 1,000,000
입출력 예
| numbers | result |
|---|---|
| [2, 3, 3, 5] | [3, 5, 5, -1] |
| [9, 1, 5, 3, 6, 2] | [-1, 5, 6, 6, -1, -1] |
문제풀이
1차원 적으로 보았을 때, 맨 첫 원소부터 순서대로 자신보다 큰 수를 찾으면 된다. 이런 풀이는 매우 간단하지만, 이중 for문을 사용하기 때문에 시간복잡도가 O(N^2)이 된다. 제한 사항을 확인해보면, 배열의 최대길이가 100만으로 매우 크기 때문에, 이러한 방법으로는 시간초과가 나게 된다. 따라서 다른 방법을 생각해 보아야 한다.
여러 풀이가 있겠지만, 나같은 경우는 Priority Queue (pq)를 이용하였다. pq에 numbers배열의 index와 원소 값을 넣은 뒤, pq를 원소 값을 기준으로 오름차순으로 정렬이 되게 한다. numbers를 for문으로 탐색하면서, pq에 값을 담은 뒤 pq의 맨 첫번째 (peek값)을 현재의 원소값과 비교하면서, 원소값이 pq의 peek값보다 작아질 때 까지 poll 하게끔 반복문을 돌렸다.
코드
Java
import java.util.*;
class Solution {
public int[] solution(int[] numbers) {
int[] answer = new int[numbers.length];
Arrays.fill(answer, -1);
PriorityQueue<int[]> queue = new PriorityQueue<>(
(o1, o2) -> o1[1] - o2[1]
);
for (int i = 0; i < numbers.length; i++) {
int temp = numbers[i];
while (!queue.isEmpty()) {
if (queue.peek()[1] >= temp) break;
answer[queue.poll()[0]] = temp;
}
queue.add(new int[] {i, temp});
}
return answer;
}
}Python
from queue import PriorityQueue
def solution(numbers):
answer = [-1 for _ in range(len(numbers))]
pq = PriorityQueue()
for i, num in enumerate(numbers):
pq.put((num, i))
while pq.queue[0][0] < num:
answer[pq.queue[0][1]] = num
pq.get()
return answer