오아시스 재결합

백준 3015번: 오아시스 재결합

아이디어

이 문제의 핵심은 "같은 키를 가진 사람들을 어떻게 처리할 것인가?" 이다. 스택에 push할 때 같은 키를 가진 사람이 몇명 연속돼있는지 함께 기록한다.

1. 스택의 top과 현재 입력값을 비교한다.
2. 현재 입력값과 같은 경우 스택의 top에 함께 저장된 연속해서 서있는 같은 사람의 수를 업데이트한다.
3. 현재 입력값이 더 작은 경우 혹은 스택이 비어있는 경우에는 그냥 push한다.
4. 현재 입력값이 더 큰 경우 top에 현재 입력값보다 큰 사람이 오거나 스택이 비어있을 때까지 pop한다. 이 때 pop되는 원소는 서로 볼 수 있는 쌍을 계산하고 pop한다. 계산은 다음과 같은 방식으로 한다.

코드

#include <bits/stdc++.h>
using namespace std;

int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    stack<pair<int, int>> s;
    int n;
    long long sz, ans = 0;
    cin >> n;
    
    while (n--) {
        int x;
        cin >> x;
        pair<int, int> p = {x, 1};
        while (!s.empty() && s.top().first < x) {
            sz = s.top().second;
            s.pop();
            if (!s.empty()) {
                ans += sz;
            }
            ans += sz;
            ans += sz*(sz-1)/2;
        }
        if (!s.empty() && s.top().first == x) {
            pair<int, int> p = {x, s.top().second+1};
            s.pop();
            s.push(p);            
        }
        else {
            s.push(p);
        }
    }
    while (!s.empty()) {
        sz = s.top().second;
        s.pop();
        if (!s.empty()) {
            ans += sz;
        }
        ans += sz*(sz-1)/2;
    }
    cout << ans;
    return 0;
}

여담

어제 하루종일 헤딩한 문제인데 오늘 풀려서 기분이 상당히 좋다. pair랑 2칸짜리 array 둘 다 써봤는데 시간차이가 별로 안난다.