[백준] 10828번 - 스택 (Python)

제출한 코드

n = int(input())
list = [list(input().split()) for _ in range(n)]

stack = []
for i in list:
  if i[0] == 'push':
    stack.append(int(i[1]))
  if i[0] == 'pop':
    if len(stack) == 0:
      print(-1)
    else:
      print(stack.pop())
  if i[0] == 'size':
    print(len(stack))
  if i[0] == 'empty':
    print(1) if len(stack) == 0 else print(0)
  if i[0] == 'top':
    if len(stack) == 0:
      print(-1)
    else:
      top = stack.pop()
      print(top)
      stack.append(top)

문제가 어렵지 않아서 단순 if ... else 문으로 구현했는데 시간 제한이 0.5초라 그런지 시간 초과로 틀렸다 😭 시간 복잡도를 고려해서 효율적으로 문제 푸는 게 아직은 익숙하지가 않다.

질문 게시판을 검색해보고 입출력 가속을 사용해야 한다는 것을 알게 되었다!

수정한 코드

import sys

n = int(sys.stdin.readline())
list = [list(sys.stdin.readline().split()) for _ in range(n)]

stack = []
for i in list:
  if i[0] == 'push':
    stack.append(int(i[1]))
  if i[0] == 'pop':
    if len(stack) == 0:
      print(-1)
    else:
      print(stack.pop())
  if i[0] == 'size':
    print(len(stack))
  if i[0] == 'empty':
    print(1) if len(stack) == 0 else print(0)
  if i[0] == 'top':
    if len(stack) == 0:
      print(-1)
    else:
      top = stack.pop()
      print(top)
      stack.append(top)

이 글에서 소개하는 더 간단하게 입출력 가속하는 방식을 써보고 싶었는데 왜인지 네임 에러가 계속 발생해서 input 자리에 sys.stdin.readline() 을 대신 써주는 식으로만 수정해봤다.

입출력 가속도 필요할 때 잘 사용할 수 있도록 유의하고 있어야겠다.

문제 출처
https://www.acmicpc.net/problem/10828

참고
https://www.acmicpc.net/board/view/67337
https://covenant.tistory.com/141