문제 바로가기
문제:
크기가 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)을 공백으로 구분해 출력한다.
풀이:
import sys
sys.stdin = open('/Users/song/Desktop/Python/Python/h.txt', 'r')
N = int(sys.stdin.readline().rstrip())
Numbers = list(map(int,sys.stdin.readline().split(" ")))
stk = [] #
res = [-1] * N # 오큰수를 저장할 리스트
for i in range(N):
while stk and Numbers[stk[-1]] < Numbers[i]:
idx = stk.pop()
res[idx] = Numbers[i]
stk.append(i)
print(*res)
해당 부분은 스택(stk)이 비어있지 않고, 현재 탐색하고 있는 숫자(Numbers[i])가 스택의 가장 위에 있는 숫자보다 큰 경우이다.
여기서 스택(stk)은 수열 Numbers의 인덱스를 저장하는 용도로 사용한다.
스택에는 현재까지 탐색한 숫자들의 인덱스가 차례로 저장되어 있다.
스택의 가장 위에 있는 인덱스에 해당하는 숫자가 현재 탐색하는 숫자보다 작다면, 그 숫자는 오큰수이다. 즉, Numbers[i]가 스택의 가장 위에 있는 숫자보다 크다는 것은 수열에서 Numbers[i]의 오큰수를 찾았다는 것이다.
그렇다면 오큰수를 찾은 인덱스에 해당하는 결과 리스트(res)의 값을 Numbers[i]로 업데이트해주면 된다.
그리고 while 루프가 있는 이유는 스택에 있는 숫자들 중에서 Numbers[i]보다 작은 모든 수들을 오큰수로 업데이트하기 위해서이다.
스택에 있는 숫자들은 현재까지 탐색한 숫자들 중에서 아직 오큰수를 찾지 못한 숫자들의 인덱스이다.
따라서 현재 탐색하는 숫자(Numbers[i])가 스택의 가장 위에 있는 숫자보다 크다면, 해당 숫자들은 모두 오큰수가 될 수 있으므로 while 루프를 통해 업데이트를 수행한다.
댓글