[PS] BOJ 3015 / 오아시스 재결합

[PS] BOJ 3015 / 오아시스 재결합
Thumbnail: Photo by Austin Neill (Unsplash)
문제 링크: 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번 보는 것과 다를 게 없습니다. 경우를 나눠서 생각해 봅시다.

  1. 새로 읽은 사람의 키가 이전 사람의 키보다 큰 경우
    1. 이전 사람들 중 현재 사람보다 키가 작은 경우는 이후의 사람들과 쌍을 이루지 못합니다. 따라서, 기억할 필요가 없습니다.
  2. 새로 읽은 사람의 키가 이전 사람과 같은 경우
    1. 만약 키가 같은 사람들이 연속되어 있는 경우, 이들은 모두 이전 / 이후의 사람들과 동일한 구성으로 쌍을 이룰 수 있습니다. 따라서, 묶어서 생각할 수 있습니다.
  3. 새로 읽은 사람의 키가 이전 사람보다 작은 경우
    1. 이전 사람들 모두가 다음 사람과 쌍을 이룰 가능성이 있습니다.

결국 우리가 배열을 읽을 동안 기억해야 하는 것은 "이번에 읽은 사람보다 키가 큰 사람들" 입니다. 또, 사람들의 키를 기억할 때, 같은 키를 가지는 사람들이 연속된다면 묶어서 생각할 수 있으니 (키, 사람 수) 형태로 저장해야 합니다.

구현

위 고찰을 토대로 생각해 볼 때, 스택에 정보를 어떻게 채우고 빼야 하는지 정리해 봅시다.

입력 배열을 순회하며, 현재 읽어온 사람의 키를 cur_height라고 합시다.

  1. 스택의 제일 위에 있는 키가 cur_height보다 작을 동안
    1. 이후의 사람들과 쌍을 이루지 못하므로, 더 이상 기억할 필요가 없습니다.
      따라서, 스택에서 제거해줍니다.
    2. 하지만, 이 사람은 cur_height와 인접해 있으니 쌍을 이룰 수 있습니다.
      이 키를 가진 사람들의 수만큼 쌍을 더해줍니다.
    3. 이 과정을 스택의 위에 있는 키가 cur_height보다 같거나 커질 때 까지 반복합니다. (스택은 위로 갈 수록 작은 키가 저장되는 상황입니다.)
  2. 스택이 비어 있다면
    1. 스택에 (cur_height, 1)을 저장합니다.
    2. 마주볼 수 있는 사람이 없으므로 쌍은 늘어나지 않습니다.
  3. 스택의 제일 위에 있는 키가 cur_height와 같을 때
    1. 또, 스택의 제일 위에 있는 키의 사람들과 쌍을 이룰 수 있습니다.
    2. 스택의 제일 위에 있는 정보 (키, 사람 수)에서 사람 수만 늘려주면 됩니다
  4. 스택의 제일 위에 있는 키가 cur_height보다 작을 때
    1. 스택에 (cur_height, 1)을 저장합니다.
    2. 쌍은 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