[Python] 백준 #1874 스택 수열 - stack

문제

1874번: 스택 수열

풀이

스택과 인덱스를 이용해서 풀었다. 

N = int(input()) # 총 갯수 
li = [] # 인풋값을 순서대로 담을 배열
stack = [] # 스택으로 활용할 배열
ans = [] # push/pop을 기록할 배열

# N개의 인풋을 li에 저장하기
for _ in range(N):
    li.append(int(input()))

# 인덱스로 현재 숫자 관리
index = 1

for i in range(N):
    if li[i] >= index:
        while index <= li[i]:
            stack.append(index)
            ans.append("+")
            index += 1
        stack.pop()
        ans.append("-")
    else:
        if len(stack) != 0 and stack[-1] == li[i]:
            stack.pop()
            ans.append("-")

# 스택이 깔끔하게 비워지지 않은 경우: 불가능 (NO 출력)
if len(stack) != 0:
    print("NO")
else: # 스택이 비워짐: 진행한 push/pop 연산 순서대로 출력
    for each in ans:
        print(each)

귀찮아서 그냥 인풋값을 input()으로 받았더니 시간이 4000ms이 넘게 나왔다.

 

 

성능 개선을 위해 input()을 sys.stdin.readline()으로 바꿔줬다.

import sys

N = int(sys.stdin.readline())
li = []
stack = []
ans = []

for _ in range(N):
    li.append(int(sys.stdin.readline()))

index = 1

for i in range(N):
    if li[i] >= index:
        while index <= li[i]:
            stack.append(index)
            ans.append("+")
            index += 1
        stack.pop()
        ans.append("-")
    else:
        if len(stack) != 0 and stack[-1] == li[i]:
            stack.pop()
            ans.append("-")

if len(stack) != 0:
    print("NO")
else:
    for each in ans:
        print(each)

역시 빠르구나

메모리는 같고 시간이 거의 1/20으로 줄었다!