[Algorithm] 백준 10828

알고리즘 기초를 순차적으로 풉니다.

문제

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

명령은 총 다섯 가지이다.

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

풀이

import sys

N = int(sys.stdin.readline())
s = []

def push(num):
    s.append(num)

def pop():
    if empty():
        return -1
    return s.pop()

def size():
    return len(s)

def top():
    if empty():
        return -1
    return s[-1]
    
def empty():
    if len(s) == 0:
        return 1
    return 0

for _ in range(N):
    query = sys.stdin.readline().split()
    if 'push' in query[0]:
        push(int(query[1]))
    elif 'pop' in query[0]:
        print(pop())
    elif 'size' in query[0]:
        print(size())
    elif 'top' in query[0]:
        print(top())
    elif 'empty' in query[0]:
        print(empty())

Lesson Learn

1. 전역 변수를 함수 안에서 사용하고 싶다면

global 선언을 해서 사용하기

2. 입력을 받을때 input()은 느리다 : sys 사용하기

입력 방식은 올바르게 구현되어 있지만, 시간 초과가 발생하는 이유는 input() 함수를 사용하여 입력을 받기 때문입니다. input() 함수는 표준 입력(stdin)에서 한 줄씩 입력을 받는데, 이는 매우 느리기 때문에 반복문 안에서 여러 번 호출되면서 입력을 처리하게 되면 시간 초과가 발생할 수 있습니다.

이를 개선하기 위해서는 한 번에 여러 줄을 입력 받고 처리하는 방식으로 변경해야 합니다. 대부분의 Online Judge나 채점 시스템에서는 입력이 많을 때 이를 해결하기 위해 여러 줄을 한꺼번에 입력 받고 처리하는 방식을 권장합니다.

예를 들어, 아래와 같이 입력을 한꺼번에 받아 처리할 수 있습니다.

import sys

N = int(sys.stdin.readline())
s = []

def push(num):
    s.append(num)

def pop():
    if empty():
        return -1
    return s.pop()

def size():
    return len(s)

def top():
    if empty():
        return -1
    return s[-1]
    
def empty():
    if len(s) == 0:
        return 1
    return 0

for _ in range(N):
    query = sys.stdin.readline().split()
    if 'push' in query[0]:
        push(int(query[1]))
    elif 'pop' in query[0]:
        print(pop())
    elif 'size' in query[0]:
        print(size())
    elif 'top' in query[0]:
        print(top())
    elif 'empty' in query[0]:
        print(empty())

이 코드에서 sys.stdin.readline()을 사용하여 한 줄씩 입력 받는 것이 아니라, 여러 줄을 한꺼번에 입력 받고 개행 문자('\n')로 구분하여 처리합니다. 이렇게 하면 입력 속도가 향상되어 시간 초과를 피할 수 있습니다.