백준 # Python_17298_오큰수

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

4

3 5 2 7

예제 출력 1

5 7 7 -1

예제 입력 2

4

9 5 4 8

예제 출력 2

-1 8 8 -1

알고리즘 분류

• 자료 구조
• 스택

---

풀이 과정

stack을 사용하여 풀어야 하는 문제인데 brute force방식으로 풀었던 것이 문제였다.

그리고 시간 초과가 난 이유는 이중 for문을 사용했었기 때문이었다.

 

 

즉, O(N^2)가 걸리는 문제를 stack을 사용해 O(N)에 가깝게 만들었어야 하는 문제였던 것이다.

 

이 것을 깨닫기 전엔 코드를 계속 더 간소히 짜는 것에만 집중했기 때문에 계속 시간초과가 났었다.

 

아래는 이 것을 깨닫기 전에 만든 코드다.

from sys import stdin
N = int(stdin.readline())
numberList = list(map(int, stdin.readline().split()))

res = []
for _ in range(N):
    res.append(-1)

for i in range(0, N):
    newList = numberList[i+1:]
    for j in range(len(newList)):
        if numberList[i] < newList[j]:
            res[i] = newList[j]
            break
print(*res)

 

그리고 아래는 stack을 활용해 푼 코드다.

자세한 설명은 주석을 참고해주세요!

 

from sys import stdin

# 충 숫자의 갯수
n = int(stdin.readline())
# 숫자가 담긴 리스트
numberList = list(map(int, stdin.readline().split()))
# 검사를 시작하는 시작 인덱스를 결정해주는 stack
stack = []
# n만큼의 길이를 가지고, -1의 값으로 설정된 리스트
res = [-1 for _ in range(n)]

for i in range(len(numberList)): # 0 1 2 3
    # 만약 스텍에 값이 있고 기준이 되는 값 보다 오른쪽 값이 더 크다면
    while stack and numberList[stack[-1]] < numberList[i]:
        # res에 기준이 되는 값의 인덱스 자리에 더 큰 값을 -1 대신 집어넣고
        res[stack.pop()] = numberList[i]
    # stack에 값이 없거나 자신보다 더 큰 숫자가 오른쪽에 존재하지 않는다면
    # 다음 비교를 위해 stack에 현재 인덱스를 담는다
    # 참고로 stack에 값이 없는 경우는 첫 번째 인덱스의 경우 밖에 없다.
    stack.append(i)

print(*res)