[C/C++] 백준 17298번 - 오큰수 (스택)

 

 

<코드>

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

int N;
stack<int> s;
int ans[1000001];
int seq[1000001];

int main()
{
	cin >> N;
	// 수열 입력받기
	for (int i = 0; i < N; i++)
		cin >> seq[i];
	
	for (int i = N - 1; i >= 0; i--)
	{
		while (!s.empty() && s.top() <= seq[i])
			s.pop();

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

		s.push(seq[i]);
	}
	
	// 정답 출력
	for (int i = 0; i < N; i++)
		cout << ans[i] << " ";
}

 

풀이 방법

 

<핵심 아이디어>

만약 수열이 이런 식으로 주어졌다고 가정하겠습니다.(막대의 길이는 해당 인덱스의 수열의 크기) 그리고 스택에는 오름차순으로 스택을 쌓을 예정입니다. (top이 가장 작고 bottom이 가장 크다는 얘기)

 

그리고 위의 그림에서 초록색 막대의 오큰수를 구해야 할 때 스택에는 1~4번 빨간 막대가 오름차순으로 쌓여있습니다.(스택의 top이 1번 빨간 막대.) 초록 막대와 스택의 top과 크기를 비교하면서 스택이 top이 초록 막대보다 더 커질 때까지 스택을 pop 합니다.  

초록막대가 더 크기 때문에 pop

 

초록막대가 더 크기 때문에 pop

 

스택이 top이 더 크기 때문에 초록막대의 오큰수는 스택의 top

 

그리고 초록막대는 스택으로 push되어 스택의 top이 됩니다.

오큰수를 찾았다면 ans배열에 기록해주고 해당 수열 값은 스택에 push 되고 다시 반복문의 맨 위로 올라갑니다.

 

예제 입력 1로 설명을 드리겠습니다.

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

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

	s.push(seq[i]);
}

for문의 i는 수열의 맨 끝부터 처음으로 이동합니다. 예제 입력 1에서 N은 4이므로 i는 3=>0으로 이동합니다.

INDEX	0	1	2	3
seq	3	5	2	7
ans

 

1. i가 3일때 스택이 비어있으므로 ans[3]은 -1로 기록하고 seq[3]를 스택에 push 해줍니다. (스택 = {7})

INDEX	0	1	2	3
seq	3	5	2	7
ans				-1

 

2. i가 2일때 스택은 비어있지 않으므로 while문에 진입해서 스택의 top과 현재 수열(seq[2])과 비교합니다. 스택의 top이 현재 수열 값보다 커지거나 스택이 비어있게 될 때까지 스택을 계속 pop 해줍니다. seq[2] = 2이므로 스택의 top이 더 크므로 ans[2]에는 스택의 top인 7을 기록하고 스택에 seq[2]를 push 합니다.  (스택 = {7, 2}) 

INDEX	0	1	2	3
seq	3	5	2	7
ans			7	-1

 

3. i가 1일때 스택은 비어있지 않으므로 while문에 진입합니다. 현재 스택은 {7 , 2}이고 스택의 top은 2이므로 스택을 pop 해줍니다. 스택의 top이 7이 되었으므로 seq[1] = 5보다 값이 크므로 while문을 탈출하고 ans[1] = 스택의 top인 7을 기록하고 스택에 seq[1]을 push 합니다. (스택 = {7, 5}) 

INDEX	0	1	2	3
seq	3	5	2	7
ans		7	7	-1

 

4. i가 0일때 현재 스택은 {7 , 5}이고 스택의 top은 5이므로 seq[0] = 3보다 스택의 top이 더 큽니다. 따라서 while문을 진입하지 않고 ans[0]에는 스택의 탑인 5를 기록해주고 seq[0]을 스택에 push 하고 for문을 종료합니다. (스택 = {7, 5, 3}) 

INDEX	0	1	2	3
seq	3	5	2	7
ans	5	7	7	-1

그리고 최종적으로 스택에는 top부터 bottom까지 오름차순으로 숫자가 쌓이게 됩니다.

이중 for문으로도 오큰수를 쉽게 구할 수 있지만 시간복잡도가 O(N²)이고 N의 최댓값이 100만이므로 100만*100만은 TLE가 발생하게 됩니다. 그러나 스택을 통한 풀이는 시간복잡도가 O(N)이므로 시간 초과 없이 통과가 가능합니다.

 

 

 

www.acmicpc.net/problem/17298

17298번: 오큰수