오아시스 재결합

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

만만하게 봤다가 은근히 고전했다.

일단 문제를 푸는 큰 틀은 아래와 같다.

• people[1] ~ people[n]까지의 사람들이 줄을 서 있을 때, 서로 볼 수 있는 쌍의 수를 알고 있다고하자.
• people[n+1]이 새로 추가될 때 새로 생겨나는 쌍의 수를 구하여 더해 나가자.

이제 위의 큰 틀을 채워나감에 있어 아래와 같이 2단계로 문제를 나누어 생각해보자.

• 새로 추가될 사람(people[n+1])이 볼 수도 있는 후보들을 유지하고, 사람이 추가될 때마다 후보들을 갱신하는 문제
• 후보들 중 실제로 새로 추가될 사람이 볼 수 있는 사람의 수를 헤아리는 문제

상술한 두 문제는 모두 스택을 활용해주면 간단하게 해결해 줄 수 있다.

사족을 붙이면, 이 문제는 처음보고 "음...라인스위핑인가?"라는 의심을 했었다.

이전에 풀어봤던 경험에 비추어보면 스위핑인가 싶은 문제는 의외로 간단한 자료구조(우선순위 큐라던가, 스택이라던가)로 풀리는 경우가 많다.

스택을 항상 키가 보다 작은 사람이 더 위에(top 방향) 위치하도록 내림차순으로 관리한다고 해보자.

• 자명하게, people[1]이 줄에 추가될 때에는 후보가 없으므로 새로 헤아려줄 경우의 수가 없다.
• 역시, 자명하게, people[1]은 두번째 사람이 볼 수도 있는 후보가 되므로 스택에 push한다.

이제 people[n] (1 < n <= N)이 줄에 추가되는 상황을 떠올려보자.

• 관리되고 있는 후보들 중 people[n]이 볼 수 있는 사람들은 곧 스택의 top부터 people[n]보다 키가 작거나 같은 후보들 + 스택의 top이후로 최초로 people[n]보다 큰 첫 사람이 될 것이다.(아직 설명하지 않았지만, 후보가 잘 관리되고 있다고 가정하자)
• 따라서, 이 수를 헤아리면 곧 새로 추가된 사람이 볼 수 있는 사람의 수를 알 수 있다.
• people[n]이 줄에 추가된다면 스택의 top부터 people[n]보다 작은 키를 갖는 후보들은 더 이상 people[n+1]에 의해 보여질 수 없게 될 것이다.(people[n]이 가려버릴 것이므로)
• 따라서, 가망이 없어진 후보들을 스택에서 pop하고 people[n]을 스택에 push해주면 후보를 올바르게 갱신해줄 수 있다.
• 위에서 설명한 내용은 곧 순서대로 (1)후보로부터 새로 추가될 사람이 볼 수 있는 사람을 세는 문제, (2)새로 사람이 추가될 때 후보를 갱신하는 문제의 아이디어를 보인다.

그런데, 잘 생각해보면 (1)에서 헤아려주어야 할 숫자는 people[n]보다 키가 작거나 같은 후보까지이고, (2)에서 pop을 해주어야 하는 후보들은 보다 작은 키를 갖는 후보들까지 이므로 스택을 iteration하게 되는 짜증나는 일이 발생한다.
(같은 키를 갖는 후보를 세어주어야 하니까)

따라서, 조금 더 머리를 써서 스택을 동일하게 키의 내림차순으로 관리하고 <키, 빈도>의 쌍으로 관리하고 스택 내에 있는 모든 키(height)가 유일하도록 관리해주면 보다 문제를 쉽게 풀 수 있다.

#include <iostream>
#include <cstdint>
#include <utility>
#include <stack>

using namespace std;

int main(void)
{
  // For Faster IO
  ios_base::sync_with_stdio(false);
  cout.tie(nullptr);
  cin.tie(nullptr);

  // Read Inputs & Solve
  int n;
  cin >> n;

  int64_t answer = 0;

  int height;
  stack<pair<int, int> > st;

  cin >> height;
  st.push(pair<int, int>(height, 1));

  for (int it = 1; it < n; ++it)
  {
    cin >> height;

    if (st.top().first > height)
    {
      answer += 1;
      st.push(pair<int, int>(height, 1));
      continue;
    }

    while (!st.empty() && st.top().first < height)
    {
      answer += st.top().second;
      st.pop();
    }

    if (st.empty())
    {
      st.push(pair<int, int>(height, 1));
    }
    else if (st.top().first == height)
    {
      auto prev = st.top();
      st.pop();

      if (st.empty())
      {
        answer += prev.second;
      }
      else
      {
        answer += prev.second + 1;
      }

      prev.second += 1;
      st.push(prev);
    }
    else
    {
      answer += 1;
      st.push(pair<int, int>(height, 1));
    }
  }

  cout << answer << '\n';

  return 0;
}