[PS] BOJ 3015 / 오아시스 재결합
문제 링크: https://www.acmicpc.net/problem/3015
스택을 활용해 풀 수 있는 문제입니다.
풀이
고찰
두 사람 A와 B가 서로 볼 수 있으려면, 두 사람 사이에 A 또는 B보다 키가 큰 사람이 없어야 한다.
다시 말해, 두 사람 A와 B 사이에 있는 사람들의 키가 min(A의 키, B의 키) 보다 작거나 같아야 합니다.
스택을 활용하며 입력 배열을 순회하면서 풀 수 있어보입니다. 하지만, 스택을 어떻게 활용해야 할까요?
먼저 예제 데이터에서 답을 직접 구해보며 생각해 봅시다.
예제 데이터
- $N$ = 7
- 이 예제에서는 총 10쌍의 사람들이 서로 볼 수 있습니다!
2
4
1
2
2
5
1입력 데이터 (N명의 사람의 키)
그림으로 보면 다음과 같습니다.

가능한 쌍은 다음과 같습니다.



사람들의 키를 순차적으로 읽는다고 생각할 때, 스택에는 어떤 정보를 저장해야 할까요?
간단하게 입력 배열을 한 칸 씩 읽으면서, 새로 읽은 사람의 키를 기준으로 새로 생기는 쌍을 계산한다고 생각해 봅시다. 이 때, 이전에 있던 사람들의 키를 참고해 쌍을 계산해야 합니다.
그럼 스택에 이전 사람들의 키를 다 담으면 될까요? 그렇게 쓰게 되면 그냥 배열을 2번 보는 것과 다를 게 없습니다. 경우를 나눠서 생각해 봅시다.
- 새로 읽은 사람의 키가 이전 사람의 키보다 큰 경우
- 이전 사람들 중 현재 사람보다 키가 작은 경우는 이후의 사람들과 쌍을 이루지 못합니다. 따라서, 기억할 필요가 없습니다.
- 새로 읽은 사람의 키가 이전 사람과 같은 경우
- 만약 키가 같은 사람들이 연속되어 있는 경우, 이들은 모두 이전 / 이후의 사람들과 동일한 구성으로 쌍을 이룰 수 있습니다. 따라서, 묶어서 생각할 수 있습니다.
- 새로 읽은 사람의 키가 이전 사람보다 작은 경우
- 이전 사람들 모두가 다음 사람과 쌍을 이룰 가능성이 있습니다.
결국 우리가 배열을 읽을 동안 기억해야 하는 것은 "이번에 읽은 사람보다 키가 큰 사람들" 입니다. 또, 사람들의 키를 기억할 때, 같은 키를 가지는 사람들이 연속된다면 묶어서 생각할 수 있으니 (키, 사람 수) 형태로 저장해야 합니다.
구현
위 고찰을 토대로 생각해 볼 때, 스택에 정보를 어떻게 채우고 빼야 하는지 정리해 봅시다.
입력 배열을 순회하며, 현재 읽어온 사람의 키를 cur_height라고 합시다.
- 스택의 제일 위에 있는 키가
cur_height보다 작을 동안- 이후의 사람들과 쌍을 이루지 못하므로, 더 이상 기억할 필요가 없습니다.
따라서, 스택에서 제거해줍니다. - 하지만, 이 사람은
cur_height와 인접해 있으니 쌍을 이룰 수 있습니다.
이 키를 가진 사람들의 수만큼 쌍을 더해줍니다. - 이 과정을 스택의 위에 있는 키가
cur_height보다 같거나 커질 때 까지 반복합니다. (스택은 위로 갈 수록 작은 키가 저장되는 상황입니다.)
- 이후의 사람들과 쌍을 이루지 못하므로, 더 이상 기억할 필요가 없습니다.
- 스택이 비어 있다면
- 스택에 (
cur_height, 1)을 저장합니다. - 마주볼 수 있는 사람이 없으므로 쌍은 늘어나지 않습니다.
- 스택에 (
- 스택의 제일 위에 있는 키가
cur_height와 같을 때- 또, 스택의 제일 위에 있는 키의 사람들과 쌍을 이룰 수 있습니다.
- 스택의 제일 위에 있는 정보 (키, 사람 수)에서 사람 수만 늘려주면 됩니다
- 스택의 제일 위에 있는 키가
cur_height보다 작을 때- 스택에 (
cur_height, 1)을 저장합니다. - 쌍은 1만큼만 증가합니다.
두 사람이 마주보기 위한 조건이 "두 사람 중 키가 작은 쪽보다 더 큰 키를 가진 사람이 중간에 없을 것"인데, 설령 스택의 위에 저장된 정보가 여러 사람이 연속된 경우라 해도 이 경우 인접한 사람 1명만 쌍을 이룰 수 있기 때문입니다.
- 스택에 (
전체 코드
from collections import deque
input = open(0).readline
stack = deque() # [해당 사람의 키, 키가 같은 연속되는 사람의 수]
N = int(input())
heights = tuple(int(input()) for _ in range(N))
pairs = 0
# 스택에 저장된 값들은 스택의 최상단부터 아래로 갈수록 커진다 (더 작은 값을 새로 추가하게 됨.)
for cur_height in heights:
while stack and stack[0][0] < cur_height:
# 현재 스택의 최상단에 위치한 키가 비교하는 키(cur_height)보다 작을 동안 계속 스택에서 뺀다.
# 이들은 키가 cur_height보다 작기 때문에 cur_height 다음 사람을 볼 수 없음이 자명함
x, count = stack.popleft()
pairs += count # 해당 키를 가진 인원이 연속으로 있는 수 만큼 더해준다.
if len(stack) == 0:
stack.appendleft((cur_height, 1))
elif stack[0][0] == cur_height: # 현재 스택의 최상단에 있는 사람의 키가 비교하려는 키와 같은 경우
height, count = stack.popleft()
pairs += count # 현재 스택의 위에 위치한 키의 사람들의 수 만큼 쌍이 더 생긴다.
if len(stack) > 0: # 스택에 그 다음 사람이 있는 경우, 마찬가지로 한 쌍이 더 생긴다. ( 이 경우, )
pairs += 1 # 왜 한 쌍만 더 생기냐? stack[1] or stack[0]의 나머지 사람들 >= stack[0] > cur_height가 되므로 stack[0]의 제일 끝쪽 사람 외에는 쌍을 이루지 못한다.
stack.appendleft((height, count + 1))
else: # 현재 스택의 위에 위치한 키보다 비교하는 키(cur_height)가 더 작은 경우
stack.appendleft((cur_height, 1))
pairs += 1 # 스택의 다음 값들과는 쌍을 이루지 못한다. (stack[1] > stack[0] > cur_height 이기 때문.)
print(pairs)solution.py