백준 / 3015 / 오아시스 재결합 / 파이썬, python, pypy

 

3015번: 오아시스 재결합

 

문제가 간단해보여서 쉽게 풀 수 있을것이라 생각했다.

 

처음에는 LIS와 같은 방식으로 생각했는데, 생각하다보니 stack을 사용해서 현재 사람보다 큰 사람만 stack에 보관하면 풀 수 있었다.

 

문제를 풀면서 고려한것은, 

 

입력받은 배열을 순서대로 탐색하면서

1. 현재 사람의 키가 stack.top()의 키보다 작으면 stack.pop() -> 내림차순으로 정렬된다.

2. stack이 비어있으면 현재 사람을 stack에 추가

3. stack.top이 현재 사람과 키가 같다면 (위의 조건으로 항상 stack에는 다른 사람이 존재)

    3-1. 만약  같은 사람들 끼리는 다 볼 수 있으므로 -> 정답 += 키가 같은 사람의 수

    3-2. stack에 현재 사람의 키와 같은 사람의 수+1을 넣는다

4. stack.top이 현재 사람보다 크므로, 현재 사람의 왼쪽(전부 나보다 같거나 큰 사람이므로)과 같이 볼 수 있다. -> 정답 += 1

 

import sys
HEIGHT, CNT = 0, 1
 
def solve():
    stack = []
    answer = 0
 
    for h in arr:
        # stack에 현재 사람보다 작은 사람이 있으면 pop하고 answer에 추가
        while stack and stack[-1][HEIGHT] < h:
            answer += stack.pop()[CNT]
 
        # 스택이 비었으면 현재 사람을 그냥 추가
        if not stack:
            stack.append((h, 1))
            continue
 
        # 스택이 안 비었고, top과 지금 현재 사람의 키가 같으면
        if stack[-1][HEIGHT] == h:
            cnt = stack.pop()[CNT]
            answer += cnt
 
            # 현재 같은키보다 왼쪽이 있으므로 -> top과 현재 사람이 볼 수 있으므로 1추가
            if stack:
                answer += 1
 
            # 현재 키인 사람과 같은 사람이 cnt만큼 있으므로, 키 = h, 명수 = cnt+1를 스택에 추가
            stack.append((h, cnt + 1))
 
        # 스택이 안비었고, 왼쪽 사람보다 키가 작으므로 그 사람만 볼 수 있음 -> 스택에 현재 키를 넣고, answer에 1추가(왼쪽 사람만 볼 수 있으므로)
        else:
            stack.append((h, 1))
            answer += 1
 
    return answer 
 
N = int(sys.stdin.readline())
arr = [int(sys.stdin.readline()) for _ in range(N)]
print(solve())