[백준] [파이썬] [스택] 17298번: 오큰수

문제집 추천, 이 문제집에 나온 유형들만 공부해도 코딩테스트는 거뜬: https://www.acmicpc.net/workbook/by/BaaaaaaaaaaarkingDog

 

17298번: 오큰수 문제보기

 

제한시간: 1초
메모리 제한: 512 MB

문제

크기가 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

---

문제 설명을 하면 A=[3,5,2,7] 에서
3보다 큰 수는 5,7인데 그중 가장 왼쪽에 있는 수는 5
5보다 큰 수는 7이므로 7
2보다 큰 수는 7이므로 7
7은 오른쪽에 수가 없으므로 -1
정답: 5 7 7 -1이다.

스텍 문제만 계속 푸니깐 패턴화 되는거 같다.

물론 한번에 풀지는 못했다.

---

첫 시도 그리고 틀린 내 코드

import sys 
sys.stdin=open('input.txt','r')

if __name__=="__main__":
    N=int(input())
    nums=list(map(int, input().split()))
    res=[-1 for _ in range(N)]
    stack=[]
    for i in range(N):
        cnt=1
        while stack:
            if nums[i]>stack[-1]: # 현재값이 stack 마지막값보다 크면 
                res[i-cnt]=nums[i]    # 현재 전 위치(즉 stack을 쌓았을때 그 값의 위치)에 현재 값을 넣는다
                stack.pop()            # stack 제거
                cnt+=1
            else:
                break
        stack.append(nums[i])
    print(*res)

사실 여러 테스트 셋을 만들어했는데 테스트셋에는 모두 맞는데 계속해서 틀림이 나온다. 뭐가 잘못된건지 모르겠다.
cnt를 이용해서 이전 위치를 표현하여 stack을 쌓았을때 i값을 찾는건데...

---

정답 코드

import sys 
sys.stdin=open('input.txt','r')

if __name__=="__main__":
    N=int(input())
    nums=list(map(int, input().split()))
    res=[-1 for _ in range(N)]
    stack=[]
    for i in range(N):
        while stack:
            if nums[i]>stack[-1][1]:
                res[stack[-1][0]]=nums[i]
                stack.pop()
            else:
                break
        stack.append([i,nums[i]])
    print(*res)

이전에 2493번: 탑 문제를 참고해서 유사히 풀었다.
문제가 되었던 업데이트 할 값이 들어갈 res의 위치를 stack을 쌓을때 값만 넣는게 아니라 위치도 넣어준다

---

다른 풀이

import sys 
sys.stdin=open('input.txt','r')

if __name__=="__main__":
    N=int(input())
    nums=list(map(int, input().split()))
    res=[-1]*N
    stack=[]
    for i in range(N):
        while stack:
            if nums[i]>nums[stack[-1]]:
                res[stack.pop()]=nums[i]
            else:
                break

        stack.append(i)
    print(*res)

깔끔하십니다... stack에 값을 넣는게 아니라 index를 넣어버렸다... 아주 좋은 답인거 같다.

---