문제 154539. 뒤에 있는 큰 수 찾기

문제 링크

https://school.programmers.co.kr/learn/courses/30/lessons/154539

프로그래머스

 

 

난이도 : Lv.2 

 

문제 요약 설명

정수로 이루어진 배열 numbers가 있습니다. 배열 의 각 원소들에 대해 자신보다 뒤에 있는 숫자 중에서 자신보다 크면서 가장 가까이 있는 수를 뒷 큰수라고 합니다.
정수 배열 numbers가 매개변수로 주어질 때, 모든 원소에 대한 뒷 큰수들을 차례로 담은 배열을 return 하도록 solution 함수를 완성해주세요. 단, 뒷 큰수가 존재하지 않는 원소는 -1을 담습니다.

입력

• 정수 배열 numbers

입력 제한

• 4 ≤ numbers의 길이 ≤ 1,000,000
  • 1 ≤ numbers[i] ≤ 1,000,000

입력 예제

// input
[2, 3, 3, 5]

// ans
[3, 5, 5, -1]

 

풀이 방식

해당 문제는 최적화 문제이다. numbers의 길이가 100만이나 되기 때문에 시간 복잡도가O(N^2)이면 당연하게도 시간초과가 발생한다. 그렇기 때문에 O(N)또는 O(NlogN)의 시간복잡도를 가지는 알고리즘을 구현해야한다.

우선, 어느 시간복잡도를 가지든 for문을 한 번만 돌아야겠다는 생각이 들었다. 이중 for문을 도는 순간 시간초과가 나기 때문이다. 그래서 우선순위 큐를 통해 다음과 같이 풀었다.

• 배열의 숫자값을 기준으로 내림차순 정렬을 하는 우선순위 큐 pq를 준비한다.
• 다음 수가 pq에 저장된 가장 큰 num보다 크다면 다음 수를 해당 숫자의 뒷 큰수로 취급한다.(숫자가 있었던 자리는 idx)
•  pq에 저장된 수가 다음 수보다 큰 수가 없을 때까지 반복한다. 

pq의 자료형은 구조체이며, 해당 구조체에는 위치(idx)와 숫자값(num)을 가지고 있기 때문에 위치와 숫자값 둘 다 접근할 수 있게된다. 따라서, 이전에 순회했던 위치의 뒷 큰 수가 정해지지 않았어도, 순회 도중 큰 수를 찾으면 언제든 갱신해 줄 수 있다. 

이렇게 풀게되면 pq.push()를 단일 for문에서 한 번씩 실행되기 때문에 O(NlogN)의 시간 복잡도를 가질 것으로 추측된다.

(우선순위 큐는 red-black tree어쩌구저쩌구....시간 복잡도가 O(logN) 어쩌구저쩌구...)

정답 코드 

#include <string>
#include <vector>
#include <bits/stdc++.h>

using namespace std;

struct pos_data{
    int num;
    int idx;
};

// 내림차순 정렬
bool operator<(const pos_data& a, const pos_data& b){
    return a.num > b.num;
} 

vector<int> solution(vector<int> numbers) {
    vector<int> answer;
    answer.resize(numbers.size(), -1);
    
    priority_queue<pos_data> pq;
    
    // O(NlogN)...?
    pq.push({numbers[0], 0});
    for(int i = 1; i < numbers.size(); i++){
        // 다음 수가 pq에 저장된 가장 큰 num보다 크다면 
        // 다음 수를 해당 숫자의 뒷 큰수로 취급한다.(숫자가 있었던 자리는 idx)
        // pq에 저장된 수가 다음 수보다 큰 수가 없을 때까지 반복한다. 
        while(!pq.empty() && pq.top().num < numbers[i]){
            answer[pq.top().idx] = numbers[i];
            pq.pop();
        }
        pq.push({numbers[i], i});
    }
    return answer;
}