크기가 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)을 공백으로 구분해 출력한다.
4
3 5 2 7
5 7 7 -1
import sys
class Node:
def __init__(self,num,data):
self.number = num
self.data = data
class Stack:
def __init__(self,number):
self.list = []
self.result = [-1]*number
def make_NGE(self,number,data):
new_node = Node(number,data)
while len(self.list):
top_node = self.list[-1]
if top_node.data < data:
self.result[top_node.number] = data
self.list.pop()
else:
break
self.list.append(new_node)
def main():
number = int(sys.stdin.readline().rstrip())
new_stack = Stack(number)
for num,data in enumerate(sys.stdin.readline().split()):
new_stack.make_NGE(num,int(data))
return new_stack.result
if __name__ == "__main__":
for i in main():
print(i,end=" ")
수열 크기가 최대 1,000,000이기 때문에 이중 반복문의 경우, 무조건 시간 초과 날 것이라고 생각하여 설계 시작함.
설계에 난항을 겪어, 백준 알고리즘 분류를 보고 stack을 사용하면 한 번에 할 수 있겠다고 생각하여, 구현함.
하지만, int(data),for i in main(): print(i,end=" ")등 마음에 안드는게 많았는데 너무 많은 시간을 잡아먹을 수 없어서 그대로 진행.
코드 단순화 : while 루프에서 스택이 비어있는지 검사하기 위해 while self.list: 로 표현해도 됨. (굳이 len(self.list)로 길이를 비교하지 않아도 됨)
로직 개선 아이디어 : if top_node.data < data 면 스택을 pop하고 result를 갱신, 그 후에도 계속 검사 → “내가 더 큰 수”인 현재 값으로 바꿀 수 있는 후보들을 전부 처리
import sys
class Element:
"""인덱스와 값을 묶어서 저장하는 간단한 클래스"""
def __init__(self, index, value):
self.index = index
self.value = value
class NGEStack:
"""
오큰수(NGE) 계산을 위한 스택을 캡슐화한 클래스.
- stack: 아직 오큰수가 결정되지 않은 Element들이 쌓이는 실제 스택.
- result: 최종 결과를 저장할 리스트(-1로 초기화).
"""
def __init__(self, size):
self.stack = []
self.result = [-1] * size # 모든 위치의 기본값을 -1로 초기화
def update_nge(self, index, value):
"""
새로운 값(value)을 넣으면서, 스택의 맨 위 값(들)과 비교해
자신보다 작은 값들의 오큰수를 갱신해 준 뒤 스택에 push한다.
"""
# 스택이 비어있지 않고, 스택 top의 value가 현재 value보다 작다면
# 그 top 원소의 오큰수를 '현재 value'로 결정하고 pop한다.
while self.stack and self.stack[-1].value < value:
top_element = self.stack.pop()
self.result[top_element.index] = value
# 아직 오큰수를 못 찾은(혹은 더 큰 값을 만났을 수 있으므로) 자기 자신을 push
self.stack.append(Element(index, value))
def main():
# 1. 입력 받기
n = int(sys.stdin.readline().strip())
arr = list(map(int, sys.stdin.readline().split()))
# 2. NGEStack 인스턴스 생성
nge_stack = NGEStack(n)
# 3. arr를 순회하며 오큰수 계산
for idx, val in enumerate(arr):
nge_stack.update_nge(idx, val)
# 4. 결과 반환
return nge_stack.result
if __name__ == "__main__":
result = main()
# 공백으로 구분된 형태로 출력
print(" ".join(map(str, result)))