[백준/Python] 10828 : 스택

 

 

겉보기에는 굉장히 쉬운 조건들이여서 의문이었는데 시간 제한의 함정이 있는 문제였다.

처음엔 막무가내로 조건별로 다 if문으로 처리해버리고, 입력 받는 것도 그냥 input() 으로 입력 받았었는데,

효율적 코드 작성을 위한 시행착오를 통해 sys 모듈까지 공부해보게 되었다.

 

 

첫 번째 코드

N = int(input())
stack = []
for _ in range(N):
    word = input()
    if 'push' in word:
        parts = word.split()
        stack.append(int(parts[1]))
    elif word == 'pop':
        if stack: print(stack.pop())
        else: print(-1)
    elif word == 'size':
        print(len(stack))
    elif word == 'empty':
        if stack : print(0)
        else: print(1)
    elif word == 'top':
        if stack: print(stack[-1])
        else: print(-1)

 

처음엔 그냥 위 코드처럼 막무가내로 if와 elif로 묶어버렸다.

예시 테스트는 모두 정상적으로 되었지만, 시간 초과로 실패하였음

 

 

두 번째 코드

N = int(input())
stack = []

# 명령을 함수로 정의
def push(x):
    stack.append(x)

def pop():
    print(stack.pop() if stack else -1)

def size():
    print(len(stack))

def empty():
    print(0 if stack else 1)

def top():
    print(stack[-1] if stack else -1)

# 명령어와 함수의 매핑 정의
commands = {
    'pop': pop,
    'size': size,
    'empty': empty,
    'top': top
}

for _ in range(N):
    word = input().strip()
    if word.startswith('push'):
        _, num = word.split()
        push(int(num))
    else:
        commands[word]()

 

각 명령어 별 기능을 토대로 각각에 대한 함수를 따로 만들어서 이를 for문 내에서 호출하는 형식으로 작성해보았다.

push 외의 함수를 호출하는 방식도 elif를 4번씩 작성하지 않고, 위와 같이 commands 딕셔너리를 새로 만들어서

여기에 명령어와 함수의 매핑을 정의하여 commands[word]() 한 줄로 각 입력값에 해당하는 함수를 호출하도록 하였다.

 

그런데 이렇게 작성해도 시간 초과가 발생해서 다음으론 입력값 받는 부분을 추가로 수정해보았다.

 

 

 

최종 코드

import sys
input = sys.stdin.read

# 모든 입력을 한 번에 읽어 들임
data = input().splitlines()

N = int(data[0])
stack = []

# 명령을 함수로 정의
def push(x):
    stack.append(x)

def pop():
    if stack:
        print(stack.pop())
    else:
        print(-1)

def size():
    print(len(stack))

def empty():
    if stack:
        print(0)
    else:
        print(1)

def top():
    if stack:
        print(stack[-1])
    else:
        print(-1)

# 명령어와 함수의 매핑을 정의
commands = {
    'pop': pop,
    'size': size,
    'empty': empty,
    'top': top
}

for i in range(1, N + 1):
    word = data[i]
    if word.startswith('push'):
        _, num = word.split()
        push(int(num))
    else:
        commands[word]()

 

사실 그 전부터 다른 분들이 푼 문제풀이들 볼 때 다들 파이썬 sys 모듈로 입력처리를 많이 하시길래 궁금했었는데,

왠지 해답이 될 수도 있을 것 같아 추가로 공부해서 적용해봤다.

 

전에는 N번 반복하는 for문 내에 input을 넣었으니, input 함수를 N번 호출하는 셈이었는데

input = sys.stdin.read , data = input().splitlines() 로 바꿔서 작성하니까 이 한 줄로 전체 N개의 input 값을 한 번에 받고, 이를 줄 별로 분리하여 data라는 리스트로 저장하게 되어 더 효율적이게 입력 값을 받을 수 있었다.

 

그리고 push 명령어에 대해서는 startswith 함수를 사용하여 push로 문자열이 시작할 경우, 문자열을 split 하여 뒤의 숫자만 추출해 이를 push() 함수로 전달하여 스택에 저장하도록 구현했다.

 

이렇게 수정하니까 드디어 성공!