스택

문제

정수를 저장하는 스택을 구현한 다음, 입력으로 주어지는 명령을 처리하는 프로그램을 작성하시오.

명령은 총 다섯 가지이다.

• push X: 정수 X를 스택에 넣는 연산이다.
• pop: 스택에서 가장 위에 있는 정수를 빼고, 그 수를 출력한다. 만약 스택에 들어있는 정수가 없는 경우에는 -1을 출력한다.
• size: 스택에 들어있는 정수의 개수를 출력한다.
• empty: 스택이 비어있으면 1, 아니면 0을 출력한다.
• top: 스택의 가장 위에 있는 정수를 출력한다. 만약 스택에 들어있는 정수가 없는 경우에는 -1을 출력한다.

입력

첫째 줄에 주어지는 명령의 수 N (1 ≤ N ≤ 10,000)이 주어진다. 둘째 줄부터 N개의 줄에는 명령이 하나씩 주어진다. 주어지는 정수는 1보다 크거나 같고, 100,000보다 작거나 같다. 문제에 나와있지 않은 명령이 주어지는 경우는 없다.

출력

출력해야하는 명령이 주어질 때마다, 한 줄에 하나씩 출력한다.

예제 입력 1

14
push 1
push 2
top
size
empty
pop
pop
pop
size
empty
pop
push 3
empty
top

예제 출력 1

2
2
0
2
1
-1
0
1
-1
0
3

예제 입력 2

7
pop
top
push 123
top
pop
top
pop

예제 출력 2

-1
-1
123
123
-1
-1

문제 풀이

import sys
n=int(input())
stack=[]

def empty(stack):
    if len(stack)==0:
        return 1
    else:
        return 0

for _ in range(n):
    command=sys.stdin.readline().split()
    if len(command)==2:
        stack.append(int(command[1]))
    else:
        if command[0]=='pop':
            if empty(stack)==1:
                print('-1')
            else:
                print(stack.pop())
        elif command[0]=='size':
            print(len(stack))
        elif command[0]=='empty':
            print(empty(stack))
        elif command[0]=='top':
            if empty(stack)==1:
                print('-1')
            else:
                print(stack[-1])

간단한 문제라고 생각해서 쉽게 풀다가 시간초과로 고생했던 문제이다.

처음에 input()으로 값을 받았다가 시간초과가 자꾸 발생해서 pop을 없애보이도 하고 함수를 새로 생성해보기도 하면서 시간을 줄여보려고 했으나 입력 개수가 많아졌을 때 대비가 안됐기 때문에 다른 방식으로 접근해야했다.

입력을 받는 데에는 두 가지 방법이 있다.

input으로 입력 받기

단순하게 input() 함수를 이용하여 입력을 받는 방법이다.

n=int(input()

a,b=map(int,input().split())

lst=list(map(int,input().split()))

위의 코드처럼 다양한 입력을 받을 수 있다.

하지만 입력의 개수가 많아지면 시간초과가 발생할 수 있기 때문에 input 대신 다른 입력 방식을 알아두면 좋다.

sys.stdin.readline으로 입력 받기

import sys

n=int(sys.stdin.readline())

a,b=map(int,sys.stdin.readline())

위의 코드처럼 input 대신에 sys.stdin.readline을 쓰면 된다.

혹은

from sys import stdin

n=int(stdin.readline())

과 같이 사용할수도 있다.

대신 sys.stdin.readline()으로 받으면 문자열로 받게 되고 끝에 \n 개행문자가 포함되니

sys.stdin.readline().strip()

으로 공백을 제거 할 수 있다.