백준 문제 중 17298번, 스택, ‘오큰수’
백준 문제 중 17298번
https://www.acmicpc.net/problem/17298
문제
크기가 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)을 공백으로 구분해 출력한다
풀이
Try1
from collections import deque
n = int(input())
# 입력을 [:n]로 n개제한
nums = map(int,input().split())
# deque 자료형 선언
nums = deque(nums)
ans=[]
# nums가 존재하는 동안 반복
while nums:
# 가장 왼쪽 팝
x = nums.popleft()
for num in nums:
# num이 x보다 크고 가장 왼쪽인 경우
if x<num:
ans.append(num)
break
# nums안에 x보다 큰 수가 없을 경우
else:
ans.append(-1)
for a in ans: print(a,end=' ')
4
3 5 2 7
5 7 7 -1
주어지는 입력이 100만개 까지도 커질수 있는데 while과 for로 두번 스캔하는 코드라 O(N^2)이 되며 시간초과가 떴다.
Try2
단순히 for 와 while의 이중 반복문으로 접근한다면 입력값의 갯수가 커질수록 기하급수적으로 복잡도가 높아지므로 n전체를 스캔하는 반복문을 최대한 줄여야 하고 이를 위해서 스택에 저장해가며 스캔하는 느낌까지는 알았는데 코드의 구성을 혼자 해결하기가 쉽지가 않아 검색을 통해 참고했다.
코드의 구성은 이렇다.
- nums의 원소를 차례로 스캔한다.
- 만약 스택 stk이 비어 있다면 현재의 원소를 넣는다.
- 다음 원소를 스캔할때 스택이 비어있지 않고, 스택의 가장 상단의 원소보다 스캔중인 원소가 크다면 무조건 스캔중인 원소가 스택에 있는 원소의 오큰수가 됨을 알 수 있다.(순서대로 스캔 했으므로 스택에 있는 원소의 오른쪽에 있음은 당연하고, 스택에 있는 수보다 스캔중인 수가 큼을 이미 판별했다)
- 스캔중인 원소를 스택에 있는 원소들의 오큰수로 하여 ans를 업데이트한다.
- ans를 업데이트 할때는 인덱스로 접근하는 편이 좋으므로 스택에 넣을때 원소의 인덱스를 같이 넣는다.[nums[i],i]
from collections import defaultdict
n = int(input())
nums = [num for num in map(int,input().split())][:n]
stk = []
# 0~n-1까지 각각을 key로 -1을 value로 갖는 딕셔너리 선언
ans = defaultdict(int)
for i in range(n): ans[i]-=1
for i in range(n):
# (2) 이후 stk이 비거나 스택의 최상단값인
# stk[-1][0]와 nums[i]를 비교
while stk and stk[-1][0]<nums[i]:
#ans[스택에서 pop한 값의 인덱스] = nums[i]
ans[stk.pop()[1]] = nums[i]
# (1) 초기에는 stk이 비어있으므로 바로 스택에 추가
# (3) 다시 스택이 비었다면 스택에 추가
stk.append([nums[i],i])
print(*ans.values())
5
3 2 4 1 8
4 4 8 8 -1
Leave a comment