백준 / 17298번 오큰수 - 스택 / Python 파이썬

 

*문제 출처는 백준에 있습니다.

문제 제목: 오큰수 (스택)

문제 사이트: https://www.acmicpc.net/problem/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

 

---

나의 풀이(틀림..)

import sys
from collections import deque
input = sys.stdin.readline
n = int(input())
lst = list(map(int,input().split()))

# stack 문제
def main(lst):
    answer = deque([-1])
    num = lst[-1]

    for i in range(1,len(lst)):
        a = len(lst) - (i + 1)
        pre_num = lst[a+1]
        if lst[a] <= num:
            if lst[a] < pre_num:
                num = pre_num
            answer.appendleft(num)
        else:
            answer.appendleft(-1)
    return answer
for i in main(lst):
    print(i,end =' ')

 

간단하게 이전 수만 비교했는데 틀림 ㅋ 흑흑흑 어렵당

 

 

모범답안

# stack 문제
def main(lst, n):
    answer = [-1] * n   # 모든 값의 오큰수를 -1로 초기화
    tmp = []            # 오큰수를 아직 못 찾은 인덱스를 저장하는 스택

    for i in range(n):
        # 스택이 비어있지 않고, 현재 값이 스택 top 인덱스의 값보다 크다면
        while tmp and lst[tmp[-1]] < lst[i]:
            # 현재 값 lst[i]가 오큰수이므로 정답 배열에 넣기
            answer[tmp.pop()] = lst[i]
        
        # 현재 인덱스를 스택에 저장
        # 아직 자기보다 더 큰 수를 못 만났기 때문에 대기 상태
        tmp.append(i)

    return answer

입력: 2 5 1 3

 

i	lst[i]	stack(tmp) -> index 저장	answer
0	2	[0]	[-1, -1, -1, -1]
1	5	[]	[5, -1, -1, -1] ← 2보다 큰 수 발견
2	1	[1, 2]	[5, -1, -1, -1]
3	3	[1]	[5, -1, 3, -1] ← 1보다 큰 수 발견

결과: 5 -1 3 -1

 

---

※ 알아야 할 것

• tmp는 아직 오큰수를 못 찾은 인덱스들을 저장.
• 현재 수가 스택 안 인덱스들의 수보다 크면 → 그 수가 오큰수임 → answer[인덱스] = 현재 값
• 스택에는 인덱스를 넣고, 비교는 lst[인덱스]로!