본문 바로가기

Coding test

[백준/17298/파이썬] 오큰수 - stack


소스코드

k = int(input())
ocs = list(map(int,input().split()))
answer = [-1] * k
stack = []
for i in range (k):
    while stack and ocs[i] > ocs[stack[-1]]:
        answer[stack[-1]] = ocs[i]
        stack.pop()
    stack.append(i)
print(*answer)

알고리즘

자기보다 큰 수를 찾지 못하는 경우는 -1을 갖게 되므로 초기 세팅을 -1 * k로 해준다.

for문을 돌며 자기보다 큰 수를 발견하면 스택에 있는 인덱스를 answer[인덱스]에 그 숫자를 넣어준다.

3 5 2 7 이 입력 되었다고 생각을 해보자 

0 1 2 3

1. 3을 비교하는데 3의 인덱스인 0을 스택에 넣는다. 

stack = [0], answer = [-1,-1,-1,-1]

 

2. 5를 비교할 차례이다.

ocs[1] -> 5 > ocs[(stack[-1] -> 0)] -> 3 성립하게 된다. 그러면 while문 안으로 들어가보자

answer[stack[-1]  -> 0 ]  = ocs[1] 즉, answer[0] = 5가 된다.

stack = [1], answer = [5,-1,-1,-1]

 

3. 2를 비교해봅시다.

ocs[2] < ocs[1] 이므로 stack에 넣기만 하고 pass

stack = [1,2], answer = [5,-1,-1,-1]

 

4. 7 마지막

ocs[3]과 stack의 마지막 값인 2을 가져와 ocs[2]를 비교하게 된다.  while문으로 간다.

answer[2] = ocs[3] = 7이 된다.

stack = [1], answer = [5,-1,7-1]

 

answer[1] = ocs[3] = 7

stack = [], answer = [5,7,7,-1]

 

종료

 

stack에 인덱스를 넣어서 문제를 풀어야 했다