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

 

 

문제 링크

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

 

 

일단 해당 문제의 티어를 보지 않고 무작정 덤빈 과거의 나에게 분노를 표한다. Grrrr

 

처음 봤을 때 드는 생각은 당연히 이중 for문을 돌려서 i와 j번째 사람의 키를 비교하는 것. 답은 당연히 나오겠지만 안 봐도 시간 초과가 나올 것이다. 따라서 기각.

 

그 다음 든 생각은, 어차피 키가 작은 사람 다음에 그 사람보다 키가 큰 사람이 있다면, 그 뒤에 다른 사람은 그 사람을 절대로 볼 수 없다. 그래서 Stack에 사람들의 키를 저장하며 단순 1만큼 카운트하다, 키가 큰 사람이 들어오면 그 사람보다 키가 작은 사람을 모두 pop해서 단순 1 카운트 + 그만큼 카운트를 세는 것이다.

 

그런데 문제가 있었다. 바로 '키가 같은 사람끼리는 모두 서로 볼 수 있다'는 사실 때문에, 같은 키의 사람이 여러 명 들어오면 계산이 이상해진다. 따라서 이것을 처리할 수 있는 로직이 있어야 된다. 나는 키를 저장하는 스택과 해당 키를 가진 사람의 수를 저장하는 스택 2개를 동시에 사용하여, 동일한 키가 들어올 때만 사람 수만큼 카운트를 해준다.

 

2, 4, 1, 2, 2, 5, 1 의 예시를 들어보면

 

[2], 0쌍

[4], 2는 1명이었기 때문에 0+1 = 1쌍

[4, 1], 단순 추가, 1+1 = 2쌍

[4, 2], 1은 1명이었고, 단순 추가로 2+1+1 = 4쌍

[4, 2], 2는 1명이었고, 단순 추가로 4+1+1 = 6쌍

[5], 2는 2명이었고, 단순 추가로 6+2+1 = 9쌍

[5, 1], 단순 추가, 9+1 = 10쌍

 

실제로 돌려보면 맞다는 걸 알 수 있다.

import sys

input = sys.stdin.readline
n = int(input())
heights = [int(input()) for _ in range(n)]

stack = [] # 사람 키만 저장
count_stack = [] # 같은 키 몇 명인지 저장
ans = 0

for h in heights:
    cnt = 1  # 현재 키의 사람 수 (기본적으로 한 명)

    while stack:
        if stack[-1] < h:
            ans += count_stack[-1]
            stack.pop()
            count_stack.pop()
        elif stack[-1] == h:
            # 같은 키면 쌍을 모두 더하고 그룹화
            ans += count_stack[-1]
            cnt += count_stack[-1]
            stack.pop()
            count_stack.pop()
        else:
            # 더 큰 사람은 하나만 볼 수 있음
            ans += 1
            break
        
    stack.append(h)
    count_stack.append(cnt)

print(ans)

 

솔직히 문제를 푼 뒤 해당 문제의 티어를 보고 놀라긴 했다. 문제에 따라서 높은 티어의 문제라도 풀 수 있는 여지가 있다는 걸 알 수 있는 시간이었다.