백준 3015번_오아시스 재결합 (Python)

문제

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

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

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

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

 

코드

from collections import deque

N = int(input())
H = []
for i in range(N):
    H.append(int(input()))
stack = deque()

ans = 0

for i in range(N):
    while len(stack) and stack[-1][0] < H[i]:
        ans += stack.pop()[1]
    
    if not len(stack):
        stack.append([H[i], 1])
    else:
        if stack[-1][0] == H[i]:
            c = stack[-1][1]
            ans += c
            stack[-1][1] = c+1
            if len(stack) > 1:
                ans += 1
        else:
            stack.append([H[i], 1])
            ans += 1
print(ans)

 

해설

관객의 수가 최대 500,000명이므로 O(N) 혹은 O(NlogN)으로 풀어야 하는 문제다. O(NlogN)는 쉽게 풀이가 떠오르지 않으므로, O(N)을 먼저 살펴본다. O(N)은 배열을 한번 다 훑고 지나가면 문제가 다 풀려 있어야 한다.

 

우선, 한 사람이 어떤 사람들을 볼 수 있는지 생각해보자. 자기보다 키가 큰 사람이 사이에 있으면 그 너머의 사람을 보지 못하므로, 앞뒤로 가장 가까운 자기보다 큰 사람까지만을 볼 수가 있다. 이 두 사람을 기점으로 한 사람이 볼 수 있는 사람들의 범위를 결정할 수가 있다.(이를 가시 범위라고 하자) 그리고 가시 범위에 속한 기준사람보다 작은 사람들 중, 기준사람과 자신 사이에 있으면서 자신보다 큰 사람이 있는 사람도 역시 서로를 보지 못한다.

 

사람들의 키에 대한 배열을 앞에서부터 순서대로 검토한다. 이때 단방향으로만 탐색을 하는데, 어차피 구하는 것은 서로 보는 사람의 '쌍'의 수이므로 한 방향으로만 검토해도 모든 '쌍'을 찾기에는 충분하기 때문이다.

그럼 진행 방향으로 가시 범위가 아직 남은, 즉 자신보다 키가 큰 사람이 나오지 않은 사람들을 아직 볼 수 있는 사람이 남은 '대기 상태'로 지정하고, 새로운 배열에 순서대로 저장해두자. 이때 대기 상태는 내림차순이어야 하는데, 아니고 더 뒤에 있는 어떤 사람이 그 앞의 사람보다 키가 크게 되어 가시 범위가 결정난 사람이 계속 대기 상태가 되기 때문이다. 참고로 대기 상태 목록에 있는 사람들은 자기와 이웃한 사람들만 볼 수 있다.

 

 

배열을 탐색하면서 새로운 사람에 대해 검토한다고 하자. 우선, 대기 상태 배열의 맨 끝에 있던 사람보다 키가 작다면 그냥 대기 상태로 지정하고 배열에 넣은 후, 맨 끝의 사람과 서로 볼 수 있으므로 서로를 볼 수 있는 쌍에 1을 더해주면 된다.

 

또, 그 사람보다 앞에 있었던 사람들 중 키가 더 작은 사람들은 그 사람을 끝으로 가시 거리가 끝이 난다. 따라서 대기 상태에서 빠져나가므로 배열에서 빠져야 한다. 이때 아직 가시 거리 배열이 비지 않았다면 가시 거리의 끝점에 있는 '새로운 사람'과 서로 볼 수 있으므로 서로를 볼 수 있는 쌍에 1을 더한다. 이때, 내림차순인 대기 상태 배열의 맨 '끝'에서부터 사람들이 순서대로 빠져야 하므로 '스택' 자료구조를 사용해야 함을 알 수 있다.

 

그리고, 대기 상태 배열의 맨 끝이 새로운 사람과 키가 같을 수도 있다. 이때는 그 사람과 그 사람을 넘어서까지도 볼 수 있으므로 그냥 바로 대기 상태 배열에 넣고 끝내기에는 나중에 빠져야 할 때 계산이 복잡해진다. 따라서 조금 특별한 방법을 쓸 것인데, 대기 상태 배열에 들어가는 사람들에 대해 이때까지 자신과 키가 같은 사람들이 가시 범위에 몇명까지 확인되었는지 기입할 것이다.

그래서 원래 있던 새로운 사람과 키가 같은 사람이 자신의 앞까지 본인포함 총 c명이 가시 범위에 확인되었다면, 새로운 사람은 그 c명을 볼 수 있을 것이므로 서로를 볼 수 있는 쌍에 c를 더해주고, 자신은 가시 범위에서 c+1명을 확인했다고 표시한 후 원래 있던 사람과 대기 상태 배열에서 대체한다. 이때도 자기 바로 앞에 (이번에는 자신보다 큰) 사람이 남아 있으면, 그 사람과도 서로 볼 수 있으므로 정답에 1을 더해준다.

 

이렇게 하여 구한 정답을 출력해주면 된다.