[Python]Stack 백준 17298

문제

크기가 N인 수열 A = A1, A2, ..., AN이 있다. 수열의 각 원소 Ai에 대해서 오큰수 NGE(i)를 구하려고 한다.

 Ai의 오큰수는 오른쪽에 있으면서 Ai보다 큰 수 중에서 가장 왼쪽에 있는 수를 의미한다. 그러한 수가 없는 경우에 오큰수는 -1이다.

예를 들어, A = [3, 5, 2, 7]인 경우 NGE(1) = 5, NGE(2) = 7, NGE(3) = 7, NGE(4) = -1이다.

 A = [9, 5, 4, 8]인 경우에는 NGE(1) = -1, NGE(2) = 8, NGE(3) = 8, NGE(4) = -1이다.

입력

첫째 줄에 수열 A의 크기 N (1 ≤ N ≤ 1,000,000)이 주어진다. 둘째에 수열 A의 원소 A1, A2, ..., AN (1 ≤ Ai ≤ 1,000,000)이 주어진다.

출력

총 N개의 수 NGE(1), NGE(2), ..., NGE(N)을 공백으로 구분해 출력한다.

처음생각할수 있는건 단순 비교이다.

오른쪽으로 가면서 비교 후 큰수가 있으면 넣고 다시돌아가 다음 수에서 시작하며 끝까지 가서 없다면 -1을 출력하는 방식이다.

이런식으로 구현해 보니 시간초과가 뜬다.

O(N^2)의 시간이 걸릴것이다. 

단계별로 풀어보기 스택부분의 마지막 문제라 그런지 어려웠다. 

힌트를 조금 얻어보았다.

스택에 index를 넣어서 확인해 보는 방식이다.

스택에 큰수를 찾지 못한 인덱스를 넣어주고 큰수를 찾게되면 큰수를 찾지 못할때까지 pop을 통해 결과를 수정하는 방식이다.

예를 들어 11 7 6 10 1 이있다면 a[0],a[1]....a[4]라고 하고 결과값을 담는 result[0],result[1].....,result[4]는 모두 -1로 초기화해준상태로 시작한다.

처음에 스택에 인덱스값 0,1,2 이 차례로 들어가고(뒤에값들이 모두 앞에값보다 작기때문에) 

 a[3]에 해당하는 10을 만나게되면 10과 비교하여 작은 6에 해당하는 인덱스의 result[2]와

 이어 7에 해당하는 인덱스 result[1]을 10으로 바꿔주고 스택에서 pop해준다 .result[0] 11과의 비교했을때는 작으므로

 비교문을 빠져나와 다시 a[4]로가서 반복한다.

 (스택에 쌓인다는건 이전값보다 작다는것이기때문에 모든값을 비교할필요없이 작은값을 만나면 다음 index에 해당하는 배열순서로 가서 

다시 실행하면된다.)

이런식으로 수열이 끝나고 나면 스택에 남는 인덱스값들은 모두 초기화상태인 -1로 남게된다.

import sys

n=int(input())

input=list(map(int,sys.stdin.readline().split())) #입력받은 수열
stack=[]      #인덱스를 담아줄 스택
result=[-1 for _ in range(n)] #모두 -1로 초기화

stack.append(0) 
i=1
while stack and i<n: #수열의 수만큼 반복
    while stack and input[stack[-1]]<input[i]:# 스택의 탑에 해당하는 인덱스 값이 입력받은 수열보다 작으면
        result[stack[-1]]=input[i]#result[index]에 위치하는 값(result[stack[-1]])을 input[i]로 바꾸어준다.
        stack.pop()#스택에서 해당 index를 빼주고 다시 비교문실행

    stack.append(i) #스택에 인덱스를 넣는다
    i+=1#인덱스를 하나늘린다.

for i in range(n):
    print(result[i],end=" ")

주말동안 많이 생각했다.

정답률이 높은걸로 보아 쉬운문제인거 같은데 오래걸렸다.

아직 많이 모자란거 같다.