17298: 오큰수

https://www.acmicpc.net/problem/17298

17298번: 오큰수

설명

크기가 N인 수열 A = A1, A2, ..., AN이 있다. 수열의 각 원소 Ai에 대해서 오큰수 NGE(i)를 구하려고 한다. Ai의 오큰수는 오른쪽에 있으면서 Ai보다 큰 수 중에서 가장 왼쪽에 있는 수를 의미한다. 그러한 수가 없는 경우에 오큰수는 -1이다.

예를 들어, A = [3, 5, 2, 7]인 경우 NGE(1) = 5, NGE(2) = 7, NGE(3) = 7, NGE(4) = -1이다. A = [9, 5, 4, 8]인 경우에는 NGE(1) = -1, NGE(2) = 8, NGE(3) = 8, NGE(4) = -1이다.

핵심 아이디어

처음에 2중 for문으로 풀려 했는데, 수열의 최대 크기가 1,000,000 이므로 시간 초과 발생한다. 따라서 스택 자료구조를 활용해야 한다. 

• stack 에는 원소값이 아닌 원소의 인덱스를 넣어주는 목적으로 사용하였다. 
• 처음에는 오큰 수가 없다고 초기화 한다.
• stack에 0 index를 삽입 한다. 그리고 첫 번째 index부터 n-1까지 비교 과정을 통해서 오큰수를 대입한다.

예를들어, 3 5 2 7 이라는 수열이 있을 때, 처음 스택에 0이 들어가 있으며, A[1]과 stack[-1]의 원소를 비교한다. 그 후 stack[-1]은 0이기 때문에 A[1]과 A[0]을 비교하는 것이며, A[1]의 값이 더 크므로 answer[0]에는 A[1]의 값인 5를 넣어주는 것이다. 

그 후 0 인덱스는 오큰수를 구했기 때문에 pop을 하고, 다음으로 구해야 할 인덱스 1을 넣어주는 것이다. 주의할 점은 현재 i가 오큰 수를 구하지 못했다 하더라고 stack에 현재 index i를 추가해줘야한다. 그래야 이어서 i+1의 오큰수를 구할 수 있고, i+1의 오큰수를 구하면 자동으로 i의 오큰수도 구하게 되기 때문이다. 

코드 구현

Python

import sys
input = sys.stdin.readline

n = int(input())
A = list(map(int, input().split())) # 3 5 2 7
'''
2중 for loop -> 시간 초과 O(n^2).
따라서 stack을 이용. stack에는 원소의 index가 들어간다.
'''
# 처음에는 오큰수가 없다고 초기화
ans = [-1] * n
stack = []
# stack에는 원소가 아니라 원소의 인덱스 값이 들어간다.
stack.append(0) # index 0 삽입
for i in range(1, n): # 첫 번째 index부터 n-1까지
	while stack and A[stack[-1]] < A[i]:
		ans[stack.pop()] = A[i] # 오큰수 대입, 오른쪽이면서 가장 왼쪽
	# 오큰수 계산을 마친 후
	stack.append(i) # 다음 오큰수 계산을 위해 다음 index 삽입

print(*ans)

C++

#include <iostream>
#include <algorithm>
#include <vector>
#include <stack>
using namespace std;


int main() {
	ios::sync_with_stdio(false);
	cin.tie(0);

	int n; 
	cin >> n;
	vector<int> A; // 수열을 저장할 벡터
	for (int i = 0; i < n; i++) {
		int num;
		cin >> num;
		A.push_back(num);
	}
	// 처음에는 오큰수가 없다고 초기화
	vector<int> ans(n, -1);
	stack<int> stack; // 원소의 인덱스를 저장할 스택

	stack.push(0); // index 0 삽입
	for (int i = 1; i < n; i++) {
		while (!stack.empty() && A[stack.top()] < A[i]) {
			ans[stack.top()] = A[i];
			stack.pop(); // 다음 오큰수 계산을 위해 스택에서 제거
		}
		stack.push(i); // 다음 오큰수 계산을 위해 다음 인덱스 삽입
	}

	for (int i = 0; i < n; i++) {
		cout << ans[i] << ' '; // 공백을 두고 출력
	}
	return 0;



}