[BOJ/Python] 6198 : 옥상 정원 꾸미기

👉 문제 링크

❌❌ 첫번째 시도 (시간초과)

문제 조건인,
자신이 위치한 빌딩보다 높거나 같은 빌딩이 있으면 그 다음 모든 빌딩의 옥상은 보지 못한다.
모든 빌딩은 일렬로 서 있고 오른쪽으로만 볼 수 있다.
를 보고 이중 for문을 생각했다.

import sys
input = sys.stdin.readline

n = int(input())

stk = []

for _ in range(n):
    stk.append(int(input()))

cnt = 0
for i in range(n-1):
    for j in range(i+1, n):
        if stk[i] > stk[j]:
            cnt += 1
        else:
            break

print(cnt)

결과는 당연히 시간초과

👉 전체 코드

import sys
input = sys.stdin.readline

n = int(input())

arr = []
stk = []
cnt = 0

for _ in range(n):
    arr.append(int(input()))

for i in range(n):
    # stk에 빌딩 높이 정보가 존재하는 경우
    while stk:
        # 현재 빌딩의 높이가 top보다 크거나 같은 경우
        if stk[-1] <= arr[i]:
            # 현재 빌딩의 높이가 top보다 
            stk.pop()
        else:
            break
        
    stk.append(arr[i])
    cnt += len(stk) -1

print(cnt)

for문을 돌면서 arr[i] 번째 빌딩의 옥상을 볼 수 있는 다른 건물의 수를 더하는 방법으로 해결했다.
len(stk) -1 인 이유는 자기 자신 제외