[Python] boj no. 3015 : 오아시스 재결합

Problem

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

Idea

N명이 한 줄로 서 있을 때, N번 사람을 기준으로 마주보고 있는 사람이 N-1명 있기 위해서는 1번부터 N-1번까지의 사람이 모두 내림차순으로 서 있어야 한다. 즉, k번째 사람을 기준으로 1번부터 k-1번째 사람들 중 내림차순으로 서 있는 사람의 수가 k번째 사람이 마주보고 있는 사람의 수이다.

이를 스택을 사용하여 다음과 같은 기준으로 구현하였다.

스택이 비어있으면 입력 키 값 push 스택이 비어있지 않으면 peek 값이 입력 키 값보다 클 때까지 pop. 이 때, pop한 횟수만큼 count 스택 탐색 종료 후, 입력 키 값을 push하면서 정답에 count 값을 더함. 이 과정을 통해 스택의 top을 기준으로 데이터가 오름차순으로 정렬된다. 하지만 키가 같은 사람이 연속되어 입력되는 케이스의 경우 마주보는 사람을 정확하게 세지 못하는 일이 발생했다. 이를 (키 값, 그 키를 가진 사람들이 연속되어 서있는 수) 의 형식으로 스택에 넣어 count를 계산하는 것으로 해결하였다.

Code

import sys
from collections import deque

input = sys.stdin.readline

stack = deque()
N = int(input())
sol = 0
for i in range(N):
    high = int(input())
    cnt = 0
    same = 0
    while stack:    # high에 따른 내림차순 정렬
        if stack[len(stack)-1][0] > high:   # peek(stack) > high, 루프문 종료
            cnt += 1
            break

        temphigh, tempsame = stack.pop()
        if temphigh == high:
            same = tempsame+1

        cnt += tempsame + 1

    sol += cnt # n번째 high와 1~n-1 번째 high 사이의 순서쌍 갯수 더하기.
    stack.append((high, same))  # (높이, 순차적으로 이어진 동일 high의 개수-1)
    # print(*stack, "{}(+{})".format(sol, cnt+same))
    
print(sol)
    
print(sol)

Comment