[BOJ] 28278번: 스택 2 / Python

https://www.acmicpc.net/problem/28278

 

스택을 구현하고, 입력되는 명령에 따라 스택에 아이템 추가/삭제/조회 를 수행하는 프로그램을 구현해야 한다.

 

결과 문자열을 result에 저장하는 아래의 코드를 제출했으나 시간 초과가 발생했다.

시간 초과 발생 코드

import sys

commands = [s.rstrip("\n") for s in sys.stdin.readlines()][1:]
stack = []
result = ""
for command in commands:
    if command[0] == "1":
        num = int(command[2:])
        stack.append(num)
    elif command[0] == "2":
        result += f"{stack.pop() if stack else -1}\n"
    elif command[0] == "3":
        result += f"{len(stack)}\n"
    elif command[0] == "4":
        result += f"{0 if stack else 1}\n"
    elif command[0] == "5":
        result += f"{stack[-1] if stack else -1}\n"

sys.stdout.write(result)

조사해보니, 시간 초과가 발생하는 이유는 다음과 같다.

for 루프에서 매번 +=를 사용해 result 문자열에 append를 할 때마다 완전히 새로운 result 문자열 객체가 생성되기 때문이다.

예를 들어 아래 코드가 수행되면 무슨 일이 벌어질는지 알아보자.

result += f"{stack.pop() if stack else -1}\n"

1. f"{stack.pop() if stack else -1}\n" 부분은 새로운 임시 문자열을 생성한다.
2. result += ... 부분의 += 연산은 under the hood에서 다음 동작을 수행한다
   • 이전의 result 문자열과 새로운 임시 문자열을 결합한 새로운 문자열이 메모리 공간에 새로 할당된다.
   • 이전의 result 문자열과 임시 문자열의 메모리 공간이 해제된다.

즉 += 연산이 발생할 때마다 result 문자열 전체에 대한 메모리에 재할당 및 복사 작업이 발생하는 것이다. 이 과정은 result 문자열이 길어지면 길어질 수록 점점 더 비용이 증가한다. 

대신에 결합할 문자열들을 list에 모아두었다가, 마지막에 join() 연산을 통해 한번에 결합하는 것이 비용을 아낄 수 있다.

 

출력할 결과 문자열을 모두 result라는 list에 모아두었다가, 마지막에 "\n".join(rsult)를 출력하도록 했다.

정답 코드

import sys

commands = [s.rstrip("\n") for s in sys.stdin.readlines()][1:]
stack = []
result = []
for command in commands:
    if command[0] == "1":
        stack.append(command[2:])
    elif command[0] == "2":
        result.append(stack.pop() if stack else "-1")
    elif command[0] == "3":
        result.append(str(len(stack)))
    elif command[0] == "4":
        result.append("0" if stack else "1")
    elif command[0] == "5":
        result.append(stack[-1] if stack else "-1")

sys.stdout.write("\n".join(result))

 

추가 최적화를 진행해봤다.

sys.stdin.readlines()는  list를 새로 생성해서 반환하기 때문에, 리스트를 만드는 시간 비용과 메모리 공간 비용을 절감해봐야 겠다고 생각했다.

sys.stdin이 파일 객체처럼 동작하므로 곧 iterable이라는 점을 이용했다. 즉 for 루프나 iterator protocol을 이용 가능하다는 의미이다. (이와 관련해서는 https://comgu.tistory.com/entry/Python-Iteration의-원리Iterator-Protocol 참고)

입력의 첫줄은 필요 없기 때문에 next()로 읽은 다음 버리고, 그 다음 줄부터 for 문으로 순회하도록 했다.

추가 최적화된 정답 코드

import sys

stack = []
result = []
next(sys.stdin)
for s in sys.stdin:
    command = s.rstrip("\n")
    if command[0] == "1":
        stack.append(command[2:])
    elif command[0] == "2":
        result.append(stack.pop() if stack else "-1")
    elif command[0] == "3":
        result.append(str(len(stack)))
    elif command[0] == "4":
        result.append("0" if stack else "1")
    elif command[0] == "5":
        result.append(stack[-1] if stack else "-1")

sys.stdout.write("\n".join(result))

메모리 사용은 절반 정도, 시간도 10% 절감했다.

결과적으로 sys.stdin.readlines()로 신규 list를 생성하지 않아도 되었기 때문이다.