[백준 3015번] 오아시스 재결합 (Python3)

import sys
input = sys.stdin.readline

N,result,stack = int(input()),0,[]
for _ in range(N):
  h = int(input())
  while stack and stack[-1][0]<h:
    h1,cnt = stack.pop()
    result += cnt
  if stack:
    if stack[-1][0]==h:
      result += stack[-1][1]
      stack[-1][1] += 1
      if len(stack)>1:
        result += 1
    else:
      result += 1
      stack.append([h,1])
  else:
    stack.append([h,1])
print(result)

간단한 문제라고 생각했는데 약간 고려할 요소가 있었다.

기본적인 풀이는 히스토그램 문제 [백준 6549번] 히스토그램에서 가장 큰 직사각형 (Python3) (tistory.com) 와 유사하게 스택을 이용해서 풀어주었다.

 

알고리즘을 요약하면,

1. stack에는 키와 해당 키의 사람 명수의 리스트를 쌓는다.

2. 입력받은 키보다 작은 키를 pop 시키고, 사람 숫자만큼 결과값에 더해준다.

3. 남은 키가 입력받은 키와 같은 경우, 같은 키의 사람들만큼 결과값에 더해주고, 만약 더 큰 키의 사람이 존재할 경우, 결과값에 1을 더해준다. 그리고 현재 키의 명수를 +1 해준다.

4. 남은 키가 입력받은 키보다 큰 경우, 결과값에 +1을 해주고, [현재 키,1]의 리스트를 쌓아준다.

5. 결과를 출력한다.

 

 

중복된 키를 따로 처리하지 않고 그냥 스택에 쌓게 되면 시간초과가 발생하였다. 이유는 예상컨대 모든 입력을 같은 수로 준다면, 시간복잡도가 O(N^2)이 되기 때문일 것이다. 따라서 그냥 키에 대해서 스택을 쌓는 것이 아닌 사람의 명수로 스택을 쌓는 것이 필요한 문제였다.

 

 

오늘의 교훈) 중복을 조심하자.