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

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

문제

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

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

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

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

입력

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

둘째 줄부터 N개의 줄에는 각 사람의 키가 나노미터 단위로 주어진다. 모든 사람의 키는 ((2^{31})) 나노미터 보다 작다.

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

출력

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

 

---

풀이

문제를 푸는 시점에서 오아시스는 실제로 한참 전에 재결합했고, 때문에 백준에서는 이 문제를 재채점했다(...)

 

스택의 고급 응용문제이다.

예제 입력을 관찰해보기로 하자.

 

2 4 1 2 2 5 1

 

우리는 스택에 값을 순서대로 추가해나가면서, 방금 추가된 값이 총 몇명의 사람을 볼 수 있는지 세줄 것이다.

 

지금 막 들어온 수를 x라고 하고, 지금 stack의 top에 있는 값을 t라고 해보자.

그러면 가능한 경우의 수는 3가지이다.

 

1. x > t일때

이렇게 되면 x 이후에 들어오는 값들은 x 앞에 있는 값들을 못본다.

따라서 t가 x보다 작거나 같은 동안에는(같은 경우는 조금 있다가 더 구체적으로 생각해보자)
stack의 값들을 pop 해주고, cnt += 1을 해준다.

그리고 stack에 x를 추가해준다.

 

2. x < t일때

 

x는 t밖에 보지 못한다.

그래서 cnt += 1을 해주고 stack에 x를 추가해준다.

 

3. x == t일때

 

문제는 여기다.

아까 x > t일때 t가 x보다 작거나 같은 동안에는 stack의 값들을 pop해주면서 cnt를 1씩 늘려준다고 했었다.

그런데, pop을 해버리면 그 이후에 들어오는 값들은 이 값들을 못 보는게 되어버리는데,

볼 수도 있다!(모르겠으면 예시 입력에 대해서 이 알고리즘을 직접 손으로 따라가보라)

그러니까 우리는 정보 손실을 막기 위해서

pop을 하면서 현재 몇 개의 같은 값들이 연속해있는지 세준다음

이를 (높이, 개수)형태의 튜플 형태로 관리해줄것이다.

그리고 나중에 t보다 큰 x가 들어오게 되면

cnt에 1을 더하는 것이 아니라 pop해주면서 개수를 더해주면 될 것이다.

 

이렇게 되면, 위의 1, 2번 경우를 수정해야하는데

stack에 x를 추가할 때 단순히 x만 추가하는게 아니라

(x, 1)의 튜플을 추가해야한다.

 

코드를 첨부한다.

import sys
input = sys.stdin.readline
N = int(input())
cnt = 0
stack = []
for i in range(N):
    x = int(input())
    if stack:
        t = stack[-1][0]
        if x >= t:
            samecnt = 1
            while stack and stack[-1][0] <= x:
                if stack[-1][0] == x:
                    samecnt += stack[-1][1]
                cnt += stack.pop()[1]
            if stack:
                cnt += 1
            stack.append((x, samecnt))
        else:
            stack.append((x, 1))
            cnt += 1
    else:
        stack.append((x, 1))
print(cnt)