[백준17298] 오큰수 (C++)

📌 문제

BOJ 바로가기

🗡 풀이

배열을 전부 탐색하는 방식으로 풀면 시간 초과가 뜨는 문제이다. 그래서 스택으로 필요 없는 부분은 쳐내고 탐색하게끔 풀어야 하는데, 스택을 떠올리기가 어려운 것 같다 😥 맨날 그래프만 풀어서 그런가... 자료구조 알고리즘도 열심히 풀어봐야겠다.

비슷한 문제로 탑 과 옥상 정원 꾸미기 가 있다.

✨ 핵심 소스코드

for (int i = 1; i <= n; i++) {
	ans[i] = -1;
	// 나중에 들어오는 수가 이전 수(들)의 오큰수인 경우
	while (!st.empty() && a[st.top()] < a[i]) {
		ans[st.top()] = a[i];
		st.pop();
	}
	// 오큰수를 찾아야 하는 수
	st.push(i);
}

📍 수열의 수를 한 개씩 push하면서 ans 배열을 채워나간다.

📍 오큰수를 (아직) 찾지 못한 수는 스택에 push된다. 나중에 들어오는 수의 인덱스를 next라고 칭하자. a[next]가 st.top()보다 큰 경우, st.top()의 오큰수가 된다.

📍 오큰수를 찾았으니 ans 배열의 값을 갱신해 주고 스택에서 pop한다. next가 그 이전 수의 오큰수도 될 수 있기 때문에 while문을 돌리며 검사해 준다.

🚩 전체 소스코드

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

int main() {
	ios_base::sync_with_stdio(false); cin.tie(0); cout.tie(0);
	int n, a[1000001], ans[1000001];
	stack<int> st;
	// 입력
	cin >> n;
	for (int i = 1; i <= n; i++)
		cin >> a[i];
	// 모든 수열의 오큰수를 찾을 때까지 반복
	for (int i = 1; i <= n; i++) {
		ans[i] = -1;
		// 나중에 들어오는 수가 이전 수(들)의 오큰수인 경우
		while (!st.empty() && a[st.top()] < a[i]) {
			ans[st.top()] = a[i];
			st.pop();
		}
		// 오큰수를 찾아야 하는 수
		st.push(i);
	}
	for (int i = 1; i <= n; i++)
		cout << ans[i] << ' ';
}