[백준 - 3015] 오아시스 재결합 - C++

 

 

 

분류: 스택 /

 

문제

문제 링크

문제 설명

오아시스의 재결합 공연에 N명이 한 줄로 서서 기다리고 있다.

이 역사적인 순간을 맞이하기 위해 줄에 서서 기다리고 있던 백준이는 갑자기 자기가 볼 수 있는 사람의 수가 궁금해졌다.

두 사람 A와 B가 서로 볼 수 있으려면, 두 사람 사이에 A 또는 B보다 키가 큰 사람이 없어야 한다.

줄에 서있는 사람의 키가 주어졌을 때, 서로 볼 수 있는 쌍의 수를 구하는 프로그램을 작성하시오.

입력

첫째 줄에 줄에서 기다리고 있는 사람의 수 N이 주어진다. (1 ≤ N ≤ 500,000)

둘째 줄부터 N개의 줄에는 각 사람의 키가 나노미터 단위로 주어진다. 모든 사람의 키는 2³¹ 나노미터 보다 작다.

사람들이 서 있는 순서대로 입력이 주어진다.

출력

서로 볼 수 있는 쌍의 수를 출력한다.

 

풀이

- 기본 원리(시간 초과)

#include <iostream>
#include <stack>
using namespace std;

int main(){
  ios::sync_with_stdio(0);
  cin.tie(0);cout.tie(0);

  int n; cin >> n;
  stack<int> S;
  long long ans=0;
  while (n--) {
    int x; cin >> x;
    int cnt = 1;
    while(!S.empty() && S.top() <= x) {
      ans++;
      if (S.top() == x) cnt++;
      S.pop();
    }
    if(!S.empty()) ans++;
    while(cnt--)
      S.push(x);
  }
  cout << ans;
}

A와 B 사이에 A 혹은 B보다 큰 사람이 없으면 A와 B는 볼 수 있다.

이 문제는 줄에 새로운 사람을 추가하며 추가된 사람과 서로 볼 수 있는 사람을 세는 문제로 볼 수 있다.

새로 추가된 i번째 사람을 `x`, 그 앞에 i-1번 째 사람을 `t`라하자.

`x`가 추가 되었을 때 `t`까지 사람 중 `x`와 서로 볼 수 있는 사람의 수를 $f(x)$라 하면 다음과 같은 식으로 답을 구할 수 있다.

$$\sum_{i}f(x_i)$$

`x`를 새로 추가할 때 발생할 수 있는 경우는 크게 3가지로 분류할 수 있다.

1. `t` < `x`
2. `t` = `x`
3. `t` > `x`

 

x가 줄에 추가되었을 때

# 1번 경우: t < x

t < x

`x`는 `x앞에 x보다 큰 사람`까지 볼 수 있고 그 사이에 `t보다 크거나 같은 사람`을 볼 수 있다.

위 그림에서 `x앞에 x보다 큰 사람`은 [0]고 `t보다 작은 사람`은 [3]뿐이므로 `x`는 [0], [1], [2], [3], `t`를 볼 수 있다.

 

# 2번 경우: t = x

1번 경우와 완전히 같은 경우다. 

t = x

위 그림에서 `x앞에 x보다 큰 사람`은 [0]이고, `t보다 작은 사람`은 [3], [4]이므로 `x`는 [0], [1], [2], `t`를 볼 수 있다.

 

# 3번 경우: t > x

t > x

이 경우 역시 1, 2번 경우와 같게 생각할 수 있다.

위 그림에서 `x앞에 x보다 큰 사람`은 `t`이므로 `x`는 `t`만 볼 수 있다.

 

# 일반화

위에서 살펴본 바와 같이 세 가지 경우는 하나의 경우로 해석할 수 있다.

`x`는 `x앞에 x보다 큰 사람`까지 볼 수 있고 그 사이에 `t보다 크거나 같은 사람`을 볼 수 있다.

따라서 `x`를 추가했을 때 `x`와 서로 볼 수 있는 사람의 수를 세는 과정은 다음과 같다.

1. `x앞에 x보다 큰 사람`을 찾는다.(없다면 줄 맨 앞까지)
2. `x앞에 x보다 큰 사람`과 `x`사이에 `t보다 크거나 같은 사람`을 센다(k명이라 하자)
3. `x`가 볼 수 있는 사람은 `x앞에 x보다 큰 사람`과 `t보다 크거나 같은 사람`들 이므로 k명 혹은 k+1명이다.

`t`가 `x`보다 작은 경우 `x`뒤의 어떤 사람도 `t`를 볼 수 없으므로 `t`는  `x`이후 고려할 필요가 없다.(pop)

하지만 `t`가 `x`보다 크거나 같은 경우 `x`뒤에서도 `t`를 볼 가능성이 있기에 고려해야 한다.(스택에 저장)

 

따라서 `monotonic decreasing stack`을 만들어 `x`를 추가 하며 $f(x)$를 계산한다.

스택에 저장된 값은 모두 `t보다 크거나 같은 사람`들 이므로 `x보다 큰 사람`을 찾을 때까지 pop하며 k를 계산한다.(while 문안의 ans++)

while문이 종료된 시점에서 스택의 top은 `x앞에 x보다 큰 사람`이므로 1을 더해준다.(스택이 비어있지 않으면 ans++)

이때 `t`가 `x`와 같을 경우 `x`뒤에서 `t`를 볼 가능성이 있으므로 pop했던 사람들을 스택에 다시 넣어줘야한다.

 

- 같은 키 압축

#include <iostream>
#include <stack>
using namespace std;

struct e {int h, cnt;};

int main(){
  ios::sync_with_stdio(0);
  cin.tie(0);cout.tie(0);

  int n; cin >> n;
  stack<e> S;
  long long ans=0;
  while (n--) {
    int x; cin >> x;
    int cnt = 1;
    while(!S.empty() && S.top().h <= x) {
      ans += S.top().cnt; // t보다 크거나 같은 사람(k)
      if (S.top().h == x) cnt += S.top().cnt;
      S.pop();
    }
    if(!S.empty()) ans++; // x앞에 x보다 큰 사람
    S.push({x, cnt});
  }
  cout << ans;
}

키가 같은 사람들이 모여있는 것을 도수로 저장하면 `x`와 `t`가 같을 경우 pop했던 `t`들을 다시 push해주는 과정을 생략할 수 있다.