- 문제
https://www.acmicpc.net/problem/2493
- 문제 풀이
- N(탑의 개수) : 1 ~ 500,000
- 탑의 높이 : 1 ~ 100,000,000
탑의 개수가 최대 50만이기 때문에 O(n2)으로 풀면 TLE가 발생한다. O(nlogn) 이하의 시간 복잡도로 풀어야 한다.
그럼 어떤 알고리즘을 써야 할까?
위 그림처럼 각각의 탑은 왼쪽 방향으로 레이저를 발사한다. 따라서 송신탑보다 더 높은 탑들 중 가장 가까운 탑이 수신한다. 이를 찾기 위한 가장 편한(?) 방법은 뒤에서부터 역으로 순회하면서 하나씩 가장 가까운 탑을 찾는 것이다. 하지만 이는 최악의 경우 O(n2)의 시간 복잡도를 가진다. 아래의 그림과 같은 경우 역순으로 찾을 시 모든 탑을 순회해야 한다.
이는 스택을 사용하면 O(n)의 시간복잡도를 가진다. 함께 살펴보자.
위와 같은 예제에서 오른쪽 탑부터 역순으로 스택에 하나씩 집어넣은다고 생각해보자.
처음에는 비어있기 때문에 4가 하나 들어갈 것이다. 그다음 7을 넣는다고 생각해보자.
7을 넣을 때 스택의 top과 비교하여 새로 넣는 수가 더 크면 pop 하고 정답 배열에 새로 넣는 수의 index를 넣는다. 이 의미는 새로운 수가 수신 탑이고 빠진 수가 송신탑이라는 뜻이다. 탑을 역순으로 넣기 때문에 수신 탑은 가장 가까운 탑이라는 것이 보장된다. 만약 새로운 수가 더 작으면 어떻게 해야 할까?
이렇게 스택에 그냥 push 한 상태로 놓는다. 그리고 다시 더 높은 수인 9가 push 되면 스택의 top이 새로운 수보다 커지거나 스택이 빌 때까지 pop 하고 정답 배열에 기록한다.
마지막으로 6을 넣으면 최종적인 스택의 모습은 다음과 같다.
마지막까지 스택에 남은 수는 수신할 탑이 없는 경우라고 보면 된다. 이렇게 푸는 경우 위에서 말한 최악의 경우에도 O(n)의 시간 복잡도를 가진다. 이를 코드로 구현하면 다음과 같다.
- 코드
#include <bits/stdc++.h>
#define fastio ios::sync_with_stdio(0), cin.tie(0), cout.tie(0)
using namespace std;
int n;
int ans[500001]; // 정답 배열
struct Top {
int idx, height;
};
vector<Top> Tops;
void input() {
// 입력
cin >> n;
for (int i = 1; i <= n; i++) {
int tmp; cin >> tmp;
Tops.push_back({ i, tmp });
}
}
void solve() {
stack<Top> st;
// 스택에 넣어가면서 기록
for (int i = n - 1; i >= 0; i--) {
Top newTop = Tops[i];
if (st.empty()) {
st.push(newTop);
}
else {
while (!st.empty()) {
if (newTop.height >= st.top().height) {
ans[st.top().idx] = newTop.idx;
st.pop();
}
else {
break;
}
}
st.push(newTop);
}
}
}
int main() {
fastio;
input();
solve();
for (int i = 1; i <= n; i++)
cout << ans[i] << ' ';
}