| 시간 제한 | 메모리 제한 |
| 1초 | 256MB |
문제
오아시스의 재결합 공연에 N명이 한 줄로 서서 기다리고 있다.
이 역사적인 순간을 맞이하기 위해 줄에서서 기다리고 있던 백준이는 갑자기 자기가 볼 수 있는 사람의 수가 궁금해 졌다.
두 사람 A와 B가 서로 볼 수 있으려면, 두 사람 사이에 A 또는 B보다 키가 큰 사람이 없어야 한다.
줄에 서있는 사람의 키가 주어졌을 때, 서로 볼 수 있는 쌍의 수를 구하는 프로그램을 작성하시오.
입력
첫째 줄에 줄에서 기다리고 있는 사람의 수 N이 주어진다. (1 ≤ N ≤ 500,000)
둘째 줄부터 N개의 줄에는 각 사람의 키가 나노미터 단위로 주어진다. 모든 사람의 키는 231 나노미터 보다 작다.
사람들이 서 있는 순서대로 입력이 주어진다.
출력
서로 볼 수 있는 쌍의 수를 출력한다.
풀이
사람을 한 명 씩 순차적으로 줄을 세워보자. 이렇게 한 사람이 줄 맨 뒤에 섰을 때, 그 사람이 앞에 있는 모든 사람들과 쌍을 이룰 수 있는 횟수를 더해나가면 마지막 사람이 줄을 섰을 때 그 횟수가 답이 된다. 이 상황에서 패턴을 찾아야 한다.
우선 가장 간단한 경우다. 줄에 선 사람이 앞 사람보다 더 작으면 그 사람은 앞 사람과 쌍이 된다. 그보다 앞에 있는 사람들과는 쌍이 될 수 없다. 따라서 이 경우는 쌍이 될 수 있는 경우가 1가지만 증가한다.
다음으로 새로 합류한 사람이 그 앞 사람보다 큰 경우다. 이 사람은 그 앞에 서있는 줄들을 쭉 봤을 때 자기보다 작은 사람들과는 쌍이 될 수 있다. 하지만 그 와중에 사이에 둘 모두보다 작지 않은 사람이 하나 있으면 그 경우는 쌍이 될 수 없다. 이 부분에서 조금 복잡해지는 문제가 생기는데, 다음 경우를 생각해보면 이것을 해결할 수 있다. 기존 맨 뒷 사람보다 큰 사람이 줄에 합류하면 그 앞에 있던 사람은 이제 뒤에 올 어떤 사람과도 쌍을 이룰 수 없다. 그러므로 키가 큰 사람이 붙었을 때, 그 앞에 있던 그보다 작은 사람들은 대열에서 빼놓고 생각해도 문제가 없다.
자 그럼, 큰 사람이 붙을 때마다 그보다 작은 사람들이 대열에서 모조리 무시되기 때문에 우리가 생각할 대열은 키가 가장 큰 사람이 맨 앞에 있고 점점 작아지는 순서로 줄을 선 꼴이 된다. 대열에 새로 붙는 사람이 어떤 사람일지 생각해보자. 현재 대열의 가장 끝에 있는 사람과 키를 비교했을 때 아래와 같이
세 가지 경우가 있다. 키가 더 큰 사람, 키가 같은 사람, 키가 더 작은 사람.
키가 더 작은 사람이 온다면 이 사람은 앞 사람하고만 쌍을 이룰 수 있다. 예를 들어 현재 대열이 [5, 4, 3, 2] 라고 해보자. 새로 붙을 사람의 키가 1이라면 이 사람은 그 앞의 4번째 사람과 쌍을 이룰 수 있고 [5, 4, 3, 2, 1] 처럼 대열에 붙으면 된다.
나머지 크거나 같은 경우인데, 큰 경우는 앞에 있는 작은 사람들과 쌍을 한 번씩 이루면서 작은 사람들을 제거해야 한다. [5, 4, 3, 2, 1] 에 키가 3인 사람이 붙는다. 이 사람은 5번째 사람과 쌍을 이룰 수 있다. 4번째 사람과도 쌍을 이룰 수 있고 3번째 사람과도 쌍을 이룰 수 있다. 여기서 4, 5번째 사람은 이제 더 큰 사람이 왔으므로 이 뒤에 올 어떤 사람과도 쌍을 이룰 수 없다. 하지만 키가 같은 3번째 사람은 뒤에 더 큰 사람이 온다면 그 사람과 쌍을 이룰 수 있다. 따라서 키가 같은 사람은 대열에서 제거하지 않는다. 그럼 키가 3인 새로운 사람이 붙은 대열은 [5, 4, 3, 3] 으로 완성된다. 여기서 4번째가 된 사람은 제거된 1, 2와 쌍을 이루었고, 지금 앞에 있는 키가 같은 3번째 사람과 쌍을 이루었다. 마지막으로 순차적으로 앞을 확인했을 때 가장 가까운 큰 사람인 2번째 사람과 쌍을 이루고 그보다 앞에 있는 사람과 쌍을 이룰 수 없다.
정리해보면 이렇다. 키가 만약 같다면, 그 앞에 키가 같은 사람이 더 있을 가능성이 있다. 이 사람들은 제거되지 않았으면서도 모두 쌍을 이룰 수 있는 사람들이기 때문에 이 사람들과 모두 쌍을 이루고, 그 앞에 있을 수도 있고 없을 수도 있는 더 큰 사람과 마지막으로 쌍을 이룰 수 있다.
대열의 끝부분에서 값을 비교하고 제거, 추가해야 하기 때문에 이 문제를 풀 때 적절한 자료구조는 스택이다. 스택으로 대열을 구현하면 문제를 해결할 수 있다.
여기서 한 가지 더, 만약 같은 키의 사람이 계속 반복해서 들어오는 경우, 스택을 계속 파고들어야 하는 문제가 있다. 이 경우 시간이 너무 많이 소요될 수 있으므로, 같은 키가 연속되는 경우, 새로운 원소를 스택에 추가하지 않고, 스택에 연속된 숫자를 함께 보관해 더할 때 같은 키를 가진 사람들을 한 묶음으로 한 번에 더하고 제거하면 빠르다.
정답 코드
import sys
from collections import deque
input = sys.stdin.readline
wait = deque()
N = int(input())
wait.append([int(input()),1])
pair = 0
for n in range(1,N):
newbie = int(input())
# 키 작은 사람들 제거
while wait and wait[-1][0] < newbie:
pair += wait.pop()[1]
# 대열이 빈 경우, 새로운 사람 추가하고 다음으로
if not wait:
wait.append([newbie,1])
continue
# 마지막 사람보다 키가 작은 경우
if newbie < wait[-1][0]:
wait.append([newbie,1])
pair+=1
# 키가 같은 사람인 경우
else:
pair+=wait[-1][1]
wait[-1][1] += 1
# 대열 끝에 키 큰 사람이 있는 경우
if len(wait)>1:
pair+=1
print(pair)
