[Python/파이썬] 백준 1874번 스택 수열

1874번: 스택 수열

문제

스택 (stack)은 기본적인 자료구조 중 하나로, 컴퓨터 프로그램을 작성할 때 자주 이용되는 개념이다. 스택은 자료를 넣는 (push) 입구와 자료를 뽑는 (pop) 입구가 같아 제일 나중에 들어간 자료가 제일 먼저 나오는 (LIFO, Last in First out) 특성을 가지고 있다.

1부터 n까지의 수를 스택에 넣었다가 뽑아 늘어놓음으로써, 하나의 수열을 만들 수 있다. 이때, 스택에 push하는 순서는 반드시 오름차순을 지키도록 한다고 하자. 임의의 수열이 주어졌을 때 스택을 이용해 그 수열을 만들 수 있는지 없는지, 있다면 어떤 순서로 push와 pop 연산을 수행해야 하는지를 알아낼 수 있다. 이를 계산하는 프로그램을 작성하라.

입력

첫 줄에 n (1 ≤ n ≤ 100,000)이 주어진다. 둘째 줄부터 n개의 줄에는 수열을 이루는 1이상 n이하의 정수가 하나씩 순서대로 주어진다. 물론 같은 정수가 두 번 나오는 일은 없다.

출력

입력된 수열을 만들기 위해 필요한 연산을 한 줄에 한 개씩 출력한다. push연산은 +로, pop 연산은 -로 표현하도록 한다. 불가능한 경우 NO를 출력한다.

코드

import sys

n = int(sys.stdin.readline())
stack = []
cnt = 1
ans = []

for _ in range(n):
    num = int(sys.stdin.readline())
    while cnt <= num:
        stack.append(cnt)
        ans.append("+")
        cnt += 1
    if stack[-1] == num:
        stack.pop()
        ans.append("-")
    else:
       print("NO")
       exit()
            
for el in ans:
    print(el)

스택에 넣을 수는 cnt

현재 들어온 입력보다 cnt가 작다면 stack에 cnt값을 push해주고 cnt 값을 1 증가시킨다.

 

처음에는 이번에 입력받은 숫자 딱 하나만 스택에서 뽑아내는줄 알고 뽑아낸 숫자들, 즉 입력받은 수보다 큰 숫자들을 따로 저장해뒀다가 입력받은 숫자만 뽑고 따로 저장해둔 큰 수들을 다시 스택에 집어넣었는데 그게 아니었다. pop하면 그보다 큰 숫자들도 모두 pop돼서 사라지는 것이었다. 내가 처음에 생각한대로 코드를 만들면 중복되는 수가 있는게 아닌 이상 만들 수 없는 수열은 존재하지 않을 것이다. 문제에서는 수열을 만들 수 없는 경우는 NO를 출력하라고 하였다. 위의 코드에서는 cnt가 이번에 입력받은 숫자보다 큰 경우 스택에 입력받은 수보다 작은 수를 push할 수도 없는 상황이므로 이 수열은 이미 만들 수 없는 것이다. 그래서 NO를 출력하고 프로그램을 종료한다.

 

아래는 처음에 제출했던 코드이다. 백준에 제출하니 메모리 초과가 났다...

import sys

n = int(sys.stdin.readline())
stack = []
cnt = 1
ans = []
for _ in range(n):
    num = int(sys.stdin.readline())
    # 스택이 비어있는 경우
    if not stack:   
        if num < cnt:
            print("NO")
            exit()
        else:
            stack.append(cnt)
            ans.append("+")
            cnt += 1
            while stack[-1] != num:
                stack.append(cnt)
                ans.append("+")
                cnt += 1
            stack.pop()
            ans.append("-")
    # 스택이 비어있지 않음
    else:   
        if cnt == num:
            ans.append("+")
            ans.append("-")
            cnt += 1
        elif stack[-1] < num:   
            while stack[-1] != num:
                stack.append(cnt)
                ans.append("+")
                cnt += 1
            stack.pop()
            ans.append("-")
        else:
            while stack[-1] > num:
                ans.append("-")
                stack.pop()
            if stack[-1] == num:
                stack.pop()
                ans.append("-")
            else:
                print("NO")
                exit()
            
for el in ans:
    print(el)

우선 stack이 비었을 경우와 아닌 경우를 나눌 필요가 없었을 것 같다.

stack이 비는 경우는 처음과 정상적으로 종료되었을 때, 그리고 만들 수 없는 수열인 경우이다. 정답 코드처럼 cnt가 현재 입력받은 수보다 작을 경우 stack에 계속 push해주도록 한다면 처음에 stack에 비는 경우는 따로 생각하지 않아도 된다. 그리고 for문으로 정해진 횟수만큼 반복문을 실행하므로 정상적으로 종료되었을 때 stack이 비는 경우도 따로 고려하지 않아도 된다. 이제 남은 경우의 수는 1가지인데 cnt와 입력받은 수의 크기 비교만으로도 만들 수 없는 수열인지 판단할 수 있으므로 stack이 비었는지 여부 확인은 굳이 필요하지 않다.  그러므로 위와 같이 stack이 비었을 경우를 나누지 않아도 된다.

 

이 코드는 반복된 부분이 많고 코드가 길어지면서 가독성도 떨어질뿐만 아니라 작성자인 나조차도 하루 지나니 대충 보면 코드가 어떤 기능을 수행하는건지 잘 모르겠는 수준이다. 좀 더 코드를 깔끔하고 어떤 기능을 하는지 한 눈에 잘 보이게 만들 수 있는 연습이 많이 필요할 것 같다.