이번에는 백준 17298번 오큰수 문제를 풀어보았습니다.
이 문제는 수열의 각 원소에 대해, 오른쪽에 있으면서 자신보다 큰 수 중 가장 왼쪽에 있는 값을 찾는 문제입니다.
처음에는 오른쪽 정보를 이어서 따라가는 방식으로 접근했고, 이후에는 stack을 사용한 방식으로 다시 정리해두었습니다.
수열 A = A1, A2, ..., AN이 주어질 때, 각 원소 Ai의 오큰수 NGE(i)를 구해야 합니다.
오큰수는 다음 조건을 만족하는 값입니다.
Ai의 오른쪽에 있으면서Ai보다 크고만약 그런 수가 없다면 -1입니다.
예를 들어
[3, 5, 2, 7]에서는 결과가 [5, 7, 7, -1][9, 5, 4, 8]에서는 결과가 [-1, 8, 8, -1]이 됩니다.
이 문제는 각 원소마다 오른쪽을 끝까지 매번 다시 확인하면 시간 초과가 날 수 있습니다.
N이 최대 1,000,000이기 때문입니다.
그래서 두 가지 방식으로 풀어보았습니다.
두 방식 모두 핵심은 같습니다.
오른쪽 원소를 단순 반복으로 매번 다시 보지 않는 것입니다.
#include <bits/stdc++.h>
using namespace std;
int main() {
int N;
cin >> N;
vector<int> arr(N);
vector<pair<int, int>> ret(N);
for (int i = 0; i < N; i++) {
cin >> arr[i];
}
ret[N-1] = {arr[N-1],-1};
for (int i=N-2; i>=0; i--) {
if (arr[i] < arr[i+1]) {
ret[i].first = arr[i+1];
ret[i].second = i+1;
} else if (ret[i+1].second == -1) {
ret[i].first = arr[i];
ret[i].second = -1;
} else {
int temp_num = i+1;
while(true) {
if (temp_num == -1) {
ret[i].first = arr[i];
ret[i].second = temp_num;
break;
}
if (ret[temp_num].first > arr[i]) {
ret[i].first = ret[temp_num].first;
ret[i].second = ret[temp_num].second;
break;
} else {
temp_num = ret[temp_num].second;
}
}
}
}
for (int i = 0; i < N; i++) {
if (ret[i].second == -1) {
cout << -1 << " ";
} else {
cout << ret[i].first << " ";
}
}
return 0;
}
ret에는 각 위치의 오큰수 값과 그 오큰수의 위치를 저장한다.1로 시작한다.1로 처리한다.이 코드에서는 단순히 오큰수 값만 저장하는 것이 아니라,
해당 오큰수의 값과 위치를 같이 저장했습니다.
vector<pair<int, int>> ret(N);
의미는 다음과 같습니다.
ret[i].first : 오큰수 값ret[i].second : 그 오큰수가 있는 위치이렇게 하면 이후에 다음 후보를 따라가며 탐색할 수 있습니다.
맨 마지막 원소는 오른쪽에 아무 값도 없기 때문에,
오큰수가 존재할 수 없습니다.
그래서 처음에 이렇게 두었습니다.
ret[N-1] = {arr[N-1],-1};
즉, 위치 정보가 -1이면 오큰수가 없다는 의미로 사용했습니다.
현재 값보다 바로 오른쪽 값이 크다면,
그 값이 곧 오큰수입니다.
if (arr[i] < arr[i+1]) {
ret[i].first = arr[i+1];
ret[i].second = i+1;
}
이 경우는 가장 단순하게 바로 처리할 수 있습니다.
바로 오른쪽 값이 더 크지 않으면,
이미 구해둔 ret 정보를 이용해서 오른쪽 후보를 타고 들어갑니다.
int temp_num = i+1;
while(true) {
if (temp_num == -1) {
ret[i].first = arr[i];
ret[i].second = temp_num;
break;
}
if (ret[temp_num].first > arr[i]) {
ret[i].first = ret[temp_num].first;
ret[i].second = ret[temp_num].second;
break;
} else {
temp_num = ret[temp_num].second;
}
}
즉, 전체를 처음부터 다시 보지 않고
이미 계산된 오큰수 위치를 따라가며 탐색 수를 줄이는 방식입니다.
#include <bits/stdc++.h>
using namespace std;
int N;
int main() {
cin >> N;
vector<int> arr = vector<int>(N,0);
vector<int> ret = vector<int>(N,-1);
stack<int> stk;
for (int i = 0; i < N; i++) {
cin >> arr[i];
while (stk.size() && arr[stk.top()] < arr[i]) {
ret[stk.top()] = arr[i];
stk.pop();
}
stk.push(i);
}
for (int i : ret) {
cout << i << " ";
}
return 0;
}
ret은 기본값을 전부 1로 초기화한다.1이다.이 코드에서 스택은 값을 직접 저장하지 않고 인덱스를 저장합니다.
stack<int> stk;
스택 안에 있는 인덱스들은
아직 자기보다 큰 수를 오른쪽에서 만나지 못한 원소들입니다.
즉, 아직 답이 정해지지 않은 원소들이 쌓여 있다고 볼 수 있습니다.
수열을 하나씩 보다가 현재 값이 스택 top이 가리키는 값보다 크다면,
그 top 원소의 오큰수는 현재 값으로 확정됩니다.
while (stk.size() && arr[stk.top()] < arr[i]) {
ret[stk.top()] = arr[i];
stk.pop();
}
현재 값보다 작은 원소들이 스택 위에 연속해서 있을 수 있으므로,
하나만 처리하는 것이 아니라 while로 반복해서 처리합니다.
오큰수가 없는 경우를 위해 ret은 처음부터 전부 -1로 초기화했습니다.
vector<int> ret = vector<int>(N,-1);
이렇게 해두면 끝까지 스택에 남아 있는 원소들은 따로 처리하지 않아도 됩니다.
그 원소들은 결국 오른쪽에 더 큰 수를 찾지 못한 경우이기 때문입니다.
현재 값보다 작은 원소들의 오큰수를 다 채운 뒤에는,
현재 인덱스를 스택에 넣습니다.
stk.push(i);
즉, 현재 원소도 이후에 더 큰 수가 나올 때까지
오큰수를 아직 모르는 상태로 남겨두는 것입니다.
이 문제는 현재 원소가 이전 원소들의 답이 될 수 있다는 점에서 stack이 잘 맞습니다.
스택 위에 남아 있는 값들은 아직 뒤에서 자신보다 큰 수를 못 찾은 상태이고,
현재 들어온 수가 그보다 크다면 그 순간 답을 채울 수 있습니다.
즉, stack은
아직 답이 정해지지 않은 원소들을 관리하는 도구로 볼 수 있습니다.