백준 1874번 - 파이썬

1874번: 스택 수열

1 부터 n까지 오름차순으로 오직 stack의 push와 pop만으로 입력받은 수열을 만들 수 있는지 확인하는 문제이다.

예제를 보면서 직접 stack이 pop이 되고 push가 되는 과정을 그려보면 크게 두 가지로 나뉜다.

1. 입력받은 수열 arr1의 원소가 arr2(1~ n까지 수를 pop/push 하는데 사용될 stack)배열에 이미 존재한다면 pop
2. 존재하지 않는다면 해당숫자까지 push하고 pop을 한다.
ex) 4 3 6 8 ,,,, 이 입력된 수열이라면, arr2배열에는 아직 비었으니,
- 1부터 4까지 push하고 pop해서 4를 배출,
- 그 다음 3은 존재하니까 다시 pop,
- 그 다음 6은 없으니까 5부터 6까지 push한뒤 pop을 해서 6을 배출,,, 이런식이다.

 

처음에는 시간 초과가 나와서 sys.stdin.readlin을 사용, pypy3로 바꿔서 제출했다.

내 코드

import sys
input = sys.stdin.readline
n = int(input())
arr1 = []# 입력 받을 수열
arr2 = []# 1부터 n까지의 수를 pop/push하는데 사용될 stack
res = []# arr2를 push 와 pop한 하는 동안 pop를 한 순서대로 수열을 담을 stack(최종 arr1의 수열을 만들 수 있는지 여부 체크)
ans = []# pop/push 하는 과정을 담을 stack( + 와 - 만으로 구성)
arr1Index = 0 # 입력받은 수열을 읽는데 사용될 Index
arr2Element = 1 # 1부터 n까지의 수를 arr2에 push할 때 쓰임. ex) arr2.append(arr2Element) => 1부터 n까지 수 중 하나(오름차순으로)

for _ in range(n):
    arr1.append(int(input()))

while arr2Element <= n:
    if arr1[arr1Index] in arr2:
        res.append(arr2.pop())
        ans.append("-")
    else:
        for k in range(arr2Element, arr1[arr1Index] + 1):
            arr2.append(k)
            ans.append("+")
            arr2Element += 1
        res.append(arr2.pop())
        ans.append("-")
    arr1Index += 1

while len(arr2) > 0:
    res.append(arr2.pop())
    ans.append("-")
if res != arr1:
    print("NO")
else:
    for val in ans:
        print(val)