[백준] 17298번 오큰수 (Python 파이썬)

www.acmicpc.net/problem/17298

17298번: 오큰수

문제 설명

수열이 주어질 때, 수열의 각 원소Ai를 기준으로 오른쪽에 위치한 원소중 Ai보다 크면서 가장 왼쪽에 있는 수를 출력해야 한다. 단, 없을 경우 -1를 출력한다.

 

예를 들어, A=[3, 5, 2, 7] 이라는 수열이 있을 때 3의 오른쪽에 있으면서 3보다 큰 수 중 가장 왼쪽에 있는 원소는 5이고

5의 오른쪽에 있으면서 5보다 큰 수 중 가장 왼쪽에 있는 원소는 7이다.

그래서 [5, 7, 7, -1] 이라는 정답이 완성된다.

---

풀이 과정

처음 이 문제를 풀었을 때는, 스택을 사용하지 않고 이중for문을 이용하여 Ai를 기준으로 Ai+1 부터 n까지의 for문을 돌리며 Ai보다 큰 수를 만났을 때 break를 하는 방식으로 작성했다.

 

하지만 이 방식은 매번 N만큼의 for문을 돌면서 모든 원소들에 대해 실시해야 했기에 O(N^2)만큼의 시간복잡도가 발생했고 결국 시간초과를 초래했다.

 

그래서 stack을 이용하여 문제를 풀었다.

stack에는 원소값이 아닌 원소의 인덱스를 넣어주는 목적으로 사용하였다.

 

예를 들어, 3 5 2 7 이라는 수열이 있을 때

처음 스택에는 0이 들어가 있으며, A[1]과 A[stack[-1]의 원소를 비교한다.

stack[-1]은 0이기 때문에 A[1]과 A[0]을 비교하는 것이며, A[1]의 값이 더 크므로

answer[0]에는 A[1]의 값인 5를 넣어주게 되는 것이다. 즉 0의 오큰수를 구한 것이다.

그리고 0인덱스는 오큰수를 구했기 때문에 pop을 하고, 다음으로 구해야 할 인덱스 1을 넣어주는 것이다.

 

만약, 현재 i가 오큰수를 구하지 못했다고 하더라도 stack에 현재 인덱스i를 추가해야 한다.

그래야 이어서 i+1의 오큰수를 구할 수 있고, i+1의 오큰수를 구하면 자동으로 i의 오큰수도 구하게 되기 때문이다.

 

 

이러한 방식으로 각각의 인덱스에 대해 오큰수를 구해주면 된다.

import sys
n = int(input())
A = list(map(int, sys.stdin.readline().split()))
answer = [-1] * n
stack = []


stack.append(0)
for i in range(1, n):
    while stack and A[stack[-1]] < A[i]:
        answer[stack.pop()] = A[i]
    stack.append(i)


print(*answer)

 

https://github.com/HongEunho

전체 문제 & 코드는 위의 깃에 정리되어 있습니다.

팔로우 & 맞팔은 환영입니다 !