[백준] 17298번: 오큰수

문제

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

설명

 수열의 각 원소에 대하여 오른쪽에 있으면서 현재 원소보다 큰 수 중 가장 왼쪽에 있는 수(오큰수)를 찾는 문제이다. 이중 for문으로 풀 수 있으나 TLE이 발생한다. 더 효율적으로 찾는 방법을 구상해야 하는데, 스택을 사용한다.

코드

#include <iostream>
#include <stack>

using namespace std;

int ans[1000001];
int arr[1000001];

int main() {
	int n;
	cin >> n;

	stack<int> s;
	for (int i = 0; i < n; i++) {
		cin >> arr[i];
	}

	for (int i = n - 1;i >= 0;i--) {
		while (!s.empty() && s.top() <= arr[i])
			s.pop();

		if (s.empty())
			ans[i] = -1;
		else
			ans[i] = s.top();

		s.push(arr[i]);
	}

	for (int i = 0; i < n; i++) {
		cout << ans[i] << " ";
	}
}

---

	for (int i = n - 1;i >= 0;i--) {
		while (!s.empty() && s.top() <= arr[i])
			s.pop();

		if (s.empty())
			ans[i] = -1;
		else
			ans[i] = s.top();

		s.push(arr[i]);
	}

 문제를 풀기 위해 가장 핵심이 되는 부분이다. 배열의 마지막 원소부터 오큰수를 찾기 시작하는데, 현재 원소의 오른쪽에 있는 원소들을 스택에 저장된 상태에서 확인하기 위함이다. 예를 들어 현재 원소가 arr[i]일 때, arr[i+1]부터 arr[n-1]까지 스택에 있을 것이라고 기대할 수 있고(후술하겠지만 모두 있으리라는 보장은 없다.), 이 상태에서 오큰수를 체크하면 되기 때문이다.

 

 while문을 통해 s.top()이 현재 원소보다 작거나 같으면 s에서 제거한다. 제거된 수는 현재 원소보다 작거나 같으므로 오큰수의 조건에 의해 오큰수가 될 수가 없으며 이후 나머지 원소들의 오큰수가 될 수가 없다. "만약 이후 원소보다 제거된 수가 더 클 수 있지 않나요?"라고 한다면, 그 상황에서 오큰수는 제거된 수가 아닌 현재 원소가 더 적합할 것이다(물론 반드시 이후 원소의 오큰수가 현재 원소가 되어야 한다는 보장은 없다.). 예를 들어, [2, 4, 9, 8]라는 배열이 있다고 하자. 현재 원소가 9일 때, 스택에는 8이 존재한다. 8이 9보다 더 작으므로 8을 스택에서 제거한다. 이때, 2, 4는 8보다 작으므로 2, 4의 오큰수가 8이 될 수 있지만, 8 이전에 9가 존재하므로 2, 4의 오큰수가 결코 8이 될 수 없는 것이다. 마지막으로 정리하면, 조건에 의해 제거된 수는 결코 이후 원소의 오큰수가 될 수 없다.

 

 그 다음, 스택이 비었다면 현재 원소보다 큰 수가 존재하지 않았다는 것이고, 오큰수는 -1이 된다. 그렇지 않다면 스택의 최상단 원소가 오큰수가 된다. 마지막으로 현재 원소를 스택에 삽입하고 반복한다.