[백준/Python3] 17298 오큰수

https://www.acmicpc.net/problem/17298

---

풀이

1. 첫번째 방식으로는 기존 리스트에서 i번째 수부터 리스트의 끝까지를 슬라이싱으로 스택을 만들어 자신보다 큰 수가 나올 때 까지 pop하고 큰 수가 없을 때에는 -1을 출력하게끔 하였다. 하지만 N이 커질 때 사실상 O(n^2)의 시간복잡도가 나오는 탓인지 시간 초과가 떴다.

# 첫번째 풀이
import sys

N = int(input())
A = list(map(int, sys.stdin.readline().split()))
ans = list()

for i in range(N):
    # 자신보다 뒤에 있는 수를 모두 스택에 넣는다.
    # 편의상 스택의 끝을 앞으로 봄
    st = A[i:]
    while st:
        nge = st.pop(0)
        # 자신보다 큰 수가 나올 때까지 pop
        if nge > A[i]:
            ans.append(nge)
            break
    else:
        ans.append(-1)

print(" ".join(map(str, ans)))

2. 슬라이싱을 제거하고 한 번의 반복문 안에서 해결할 수 있도록 만들어야 시간초과가 나지 않을 것이기 때문에 방법을 바꾸었다.

# 두번째 풀이
import sys

N = int(input())
A = list(map(int, sys.stdin.readline().split()))
ans = [-1 for _ in range(N)]
# 인덱스를 저장하는 스택
stack = list()

stack.append(0)

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

print(' '.join(map(str, ans)))

기존 스택을 값을 저장하는 것이 아니라 A의 인덱스를 저장하며, 정답 리스트에는 오큰수가 없다는 가정하에 모두 -1로 초기화해둔다.
기본적으로 스택에 가장 마지막 값인 인덱스의 A의 원소와 현재 i 인덱스인 원소를 비교한다. 오큰수가 나왔을 경우 스택의 마지막 값인 인덱스를 pop하여 그 정답 위치에 오큰수를 넣는다. 이때 오큰수를 발견하지 못하더라도 다음 i+1번째 수의 오큰수를 찾으면서 i번째의 오큰수가 같이 찾을 수 있으므로 i의 인덱스는 스택에 저장한다.
스택이 있고 스택의 마지막 원소를 인덱스로 가지는 A의 원소보다 큰 수가 i에 계속 있다면 스택을 계속 pop하면서 이전에 찾지 못한 인덱스로 가지는 원소의 오큰수도 찾을 수 있다. 끝까지 오큰수가 없어 찾지 못하면 정답 리스트의 값을 변경하지 못했으므로 그대로 -1을 출력하게 된다.