백준 24511 | queuestack [파이썬 python]

24511번: queuestack

 
처음 시도한 코드 - 시간초과

from collections import deque
import sys
input = sys.stdin.readline

N = int(input())                    # queuestack을 구성하는 N개의 자료구조
A = list(map(int, input().split())) # 길이 N의 수열 A
B = list(map(int, input().split())) # 길이 N의 수열 B
M = int(input())                    # 삽입할 수열의 길이 M
C = list(map(int, input().split())) # 길이 M의 수열 C

queuestack = []
for i in range(N):
    # A[i]가 0이면 큐, 1이면 스택
    if A[i] == 0:
        queuestack.append(deque([B[i]]))
    else:
        queuestack.append([B[i]])

for m in range(M):
    item = C[m] # 삽입할 원소
    for n in range(N):
        if A[n] == 0: # 큐이면 popleft
            queuestack[n].append(item)
            item = queuestack[n].popleft()
        else: # 스택이면 pop
            queuestack[n].append(item)
            item = queuestack[n].pop()

    print(item, end=' ')

문제에서 주어진 queuestack 구성에 맞게 코드를 작성하였다.

하지만, M은 최대 1,000,000,000, N은 최대 100,000이기 때문에 이중 for문으로 하면 시간 초과가 날 수밖에 없다.

어떻게 시간 초과를 해결해야 할 지 고민하다 직접 원소 삽입 과정을 손으로 그려보니 신기한 현상을 발견할 수 있었다.
 

스택의 경우, 원소를 추가하고 pop 연산을 하는 것은 결국 자기 자신을 넣고 빼는 것이기 때문에 스택 자체에는 아무 변화가 없다.

그러나, 큐의 경우 여러개의 큐가 하나인 것처럼 작동하는 것을 알 수 있었다.

1번 큐와 4번 큐가 하나로 연결되어 있다고 생각하면 [1, 4]와 같은 형태가 된다.

이 큐에 배열 C의 원소 하나를 appendleft 해주고, pop 연산을 수행한다.

이때, pop 연산을 통해 빠져나온 원소가 문제에서 구하고자 하는 리턴값 $X_n$이 된다.

 
정답 코드

from collections import deque
import sys
input = sys.stdin.readline

N = int(input())                    # queuestack을 구성하는 N개의 자료구조
A = list(map(int, input().split())) # 길이 N의 수열 A
B = list(map(int, input().split())) # 길이 N의 수열 B
M = int(input())                    # 삽입할 수열의 길이 M
C = list(map(int, input().split())) # 길이 M의 수열 C

queue = deque([])
for i in range(N):
    if A[i] == 0: # 큐인 자료구조만 모으기
        queue.append(B[i])

# 배열 C의 원소를 1개 appendleft 할 때마다 pop 연산 수행
for i in range(M):
    queue.appendleft(C[i])
    print(queue.pop(), end= ' ')