[Python/파이썬] 백준 # 10828 스택

문제 링크: https://www.acmicpc.net/problem/10828

 

문제

 

스택에 대해 연습해 볼 수 있는 문제

주어진 다섯 가지의 명령을 구현하면 된다.

---

코드

from collections import deque

stack = deque()

n = int(input())

for _ in range(n):
    s = input()
    command = s.split()

    if command[0] == "push":
        stack.append(command[1])

    elif command[0] == "pop":
        if stack:
            print(stack.pop())
        else:
            print(-1)

    elif command[0] == "size":
        print(len(stack))

    elif command[0] == "empty":
        if stack:
            print(0)
        else:
            print(1)

    elif command[0] == "top":
        if stack:
            print(stack[-1])
        else:
            print(-1)

 

위의 코드를 제출했더니 시간 초과가 나왔다.

 

---

성능개선

from collections import deque
import sys

stack = deque()

n = int(input())
output = []

for _ in range(n):
    s = input()
    command = s.split()

    if command[0] == "push":
        stack.append(command[1])

    elif command[0] == "pop":
        if stack:
            output.append(stack.pop())
        else:
            output.append(-1)

    elif command[0] == "size":
        output.append(len(stack))

    elif command[0] == "empty":
        if stack:
            output.append(0)
        else:
            output.append(1)

    elif command[0] == "top":
        if stack:
            output.append(stack[-1])
        else:
            output.append(-1)

sys.stdout.write('\n'.join(map(str, output)))

 

처음 제출한 코드는 매번 print()를 호출하게 되어 실행 시간이 길어지는 것이라 생각해서

output이라는 리스트에 출력해야 할 값들을 저장한 후

sysstdout.write()을 사용해 한번에 출력했더니 정답 처리 되었다.