문제

https://www.acmicpc.net/problem/3015


알고리즘

stack


풀이

$N$명이 사람이 서있을 때, 서로 볼 수 있는 쌍을 구하는 문제입니다. 서로 보기 위해선 각 사람 사이에 그 사람들보다 큰 키가 존재해서는 안됩니다.

문제의 단순화를 위해 모든 키가 유니크하다고 가정하겠습니다. 파악해야할 것은 키의 구간이 내림차순으로 존재한다면 해당 구간은 추후에 매칭이 될 가능성이 있는 사람들입니다. 아래 그림은 이에 대한 예시로써, 흰색사람의 키가 내림차순인 형태입니다. 이후에 나오는 주황색 사각형과 흰색 사각형들이 매칭이 됩니다.

img

또한 빨간색 사각형처럼 자신의 왼쪽에 있는 사람이 자신보다 키가 크다면 무슨일이 있어도 주황색 사각형의 왼쪽부분과는 매칭이 될 수 없습니다. 이러한 특징을 파악해 오름차순 또는 내림차순으로 구간을 관리할 수 있는 스택을 사용하도록 합시다.

매칭될 가능성이 있는 사람들을 내림차순으로 스택으로 관리하면서 만일 스택의 top이 현재 자신보다 작다면 매칭이 이루어 질 수 있는 것입니다. top이 자신보다 작다면 계속해서 pop을하며 매칭을 해줍니다.

이 문제가 어려운 이유는 키가 유니크하지 않고 동일한 키의 사람도 존재한다는 것입니다. 연속적으로 동일한 키가 존재하는 구간이 있다면 그 구간끼리는 그들 사이에 큰 키가 존재하지 않으므로 서로 볼 수 있습니다. 이를 처리하기 위해 스택에 높이와 해당 높이의 연속적인 갯수를 저장합니다. 이렇게 저장을 하게 되면 스택에서 pop을 할 때 조합공식을 이용해 같은 키의 사람끼리 처다볼 때를 처리할 수 있습니다.

한가지 또 유의할 점은 스택에서 매칭되는 상대를 찾아 pop을 할때 자신보다 작은 키의 사람들뿐만 아니라 처음으로 등장하는 큰 키의 사람또한 현재 키의 사람과 매칭된다는 점입니다.

코드

#include <bits/stdc++.h>
#define rep(i, n) for (int i = 0; i < n; ++i)
#define REP(i, n) for (int i = 1; i <= n; ++i)
using namespace std;

const int MAXN = 5e5 + 5;
typedef pair<int, int> pii;

long long ans;
int N;
stack<pii> st;

int main() {
#ifndef ONLINE_JUDGE
    freopen("patrik.in.10", "r", stdin);
    freopen("out", "w", stdout);
#endif
    cin.tie(NULL);
    cout.tie(NULL);
    ios::sync_with_stdio(false);

    cin >> N;
    rep(i, N) {
        int h;
        cin >> h;
        while (!st.empty() && st.top().first < h) {
            int cnt = st.top().second;
            ans += cnt;
            ans += (long long)cnt * (cnt - 1) / 2;
            st.pop();
        }

        if (st.empty()) {
            st.emplace(h, 1);
        } else {
            if (st.top().first == h) {
                int cnt = st.top().second + 1;
                st.pop();

                if (!st.empty()) {
                    ans += 1;
                }
                st.emplace(h, cnt);
            } else {
                st.emplace(h, 1);
                ans += 1;
            }
        }
    }
    while (!st.empty()) {
        auto [a, b] = st.top();
        st.pop();
        ans += (long long)b * (b - 1) / 2;
    }
    cout << ans;
    return 0;
}