백준 17298 : 오큰수 (파이썬)

오큰수 성공

시간 제한	메모리 제한	제출	정답	맞힌 사람	정답 비율
1 초	512 MB	32990	11047	8074	33.192%

문제

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

---

코드

#https://teching.tistory.com/
import sys

n = int(sys.stdin.readline().rstrip())
nums = list(map(int,sys.stdin.readline().split()))
answer = [-1] * n
stack = [0]
for i in range(1,n):
    while stack and nums[stack[-1]] < nums[i]:
        answer[stack.pop()] = nums[i]
    stack.append(i)

print(*answer)

해설

처음에는 이중 for문을 통해서 앞에서 하나씩 꺼내가며 하나하나 확인하는 방법을 사용했는데 너무 오래 걸려 시간 초과가 발생했다.

그래서 스택을 이용해서 풀이하였다.

오른쪽부터 하나씩 숫자를 확인하는데

스택에는 값이 아닌 인덱스(주소)를 저장을 해주어 위치를 찾기 쉽도록 만들었고 스택의 TOP과 계속 비교해서 오큰수가 된다면 POP을 해주고 그 후에 하나씩 PUSH 해주면 된다.

 

예를 들어 3 6 2 5 7 8의 수열이 주어졌을 때 for문의 동작을 설명하자면
처음에는 stack = [0]이 들어있게 된다.

1. i=1일 때 stack의 TOP은 0이다.
   1) nums [stack [-1]] = 3,   nums[i] = 6 비교
   2) nums[i]가 더 크기 때문에 오큰수이다. 따라서 stack에서 pop을 해준다.>> stack = []
   3) pop 한 값인 0의 위치에 오큰수 6을 저장해준다. >> answer[0] = 6
   4) 현재 i를 PUSH 해준다.(아직 오큰수를 찾지 못했기 때문) >> stack = [1]
   
   
2. i=2일 때 stack의 TOP은 1이다.
   1) nums[stack[-1]] = 6, nums[i] = 2 비교
   2) nums[stack[-1]]가 더 크기 때문에 오큰수가 아니다. 따라서 더 처리해줄 것이 없다. (가장 왼쪽에서 오큰수가 안되면 그 왼쪽에 있는 오큰수를 못찾은 숫자들의 오큰수 또한 될 수 없기 때문에 더 확인해줄 필요가 없다.)
   3) i=1가 일 때와 마찬가지로 아직 i=2에 대한 오큰수를 못 찾았기 때문에 스택에 현재 i를 PUSH 해준다.
   >> stack = [1,2]
3. i=3일 때 stack의 TOP은 2이다.
   1) nums[stack[-1]] = 2, nums[i] = 5 비교
   2) nums[i]가 더 크므로 3번째 요소인 2의 오큰수는 5이다. 따라서 stack에서 pop 해준다. >> stack = [1]
   3) pop 한 값인 2의 위치에 오큰수 5를 저장해준다. >> answer[2] = 5
   4) 아직 stack에 값이 남아있으므로 추가적으로 오큰수가 될 수 있는지 확인해본다.
   5) stack의 TOP은 1이고 해당 인덱스의 값은 6이다. 5보다 크기 때문에 오큰수가 될 수 없으므로 아무 처리를 하지 않고 끝낸다.
   6) stack에 i를 PUSH 해준다. stack = [1,3]
    
4. i=4일 때 stack의 TOP은 3이다.
   1) nums[stack[-1]] = 5, nums[i] = 7 비교
   2) nums[i]가 더 크므로 3번째 요소인 5의 오큰수는 5이다. 따라서 stack에서 pop 해준다. >> stack =[1]
   3) pop 한 값인 3의 위치에 오큰수 7를 저장해준다. >> answer[3] = 7
   4) 아직 stack에 값이 남아있으므로 추가적으로 오큰수가 될 수 있는지 확인해본다.
   5) stack의 TOP은 1이고 해당 인덱스의 값은 6이다. 7보다 작기 때문에 오큰수가 될 수 있다. 따라서 stack에서 pop 해준다. >> stack = []
   6) pop 한 값인 1의 위치에 오큰수 7를 저장해준다. >> answer[1] = 7
   7) 더 이상 스택에 값이 없으므로 stack에 i를 PUSH 해주고 다음 인덱스로 넘어간다. >> stack = [4]
5. i=5일 때 stack의 TOP은 8이다.
   1) nums[stack[-1]] = 7, nums[i] =8 비교
   2) nums[i]가 더 크므로 4번째 요소인 7의 오큰수는 8이다. 따라서 stack에서 pop 해준다. >> stack = []
   3) pop 한 값인 4의 위치에 오큰수 8을 저장해준다. >> answer[4] = 8
   4) stack에 값이 없으므로 끝낸다.

이렇게 for문을 한번 돌고 나면 자연스럽게 오큰수를 찾지 못한 숫자들의 인덱스는 stack에 남아있을 것이며 해당 위치의 값은 자연스럽게 -1이 될 것이다.(처음에 -1로 초기화했기 때문)