알고리즘💻/스택&큐&덱

BOJ 17299번: 오등큰수

호프 2021. 2. 1. 20:36

www.acmicpc.net/problem/17299

 

17299번: 오등큰수

첫째 줄에 수열 A의 크기 N (1 ≤ N ≤ 1,000,000)이 주어진다. 둘째에 수열 A의 원소 A1, A2, ..., AN (1 ≤ Ai ≤ 1,000,000)이 주어진다.

www.acmicpc.net

아이디어:

주어진 수열을 돌면서 F(Ai) 배열을 채운다.

무식하게 풀어보자면.. 그냥 배열을 돌면서 다 비교할 수 있겠지만 그러면 O(n^2)이므로 시간초과가 날것이다..

 

구글링을 해보니, 스택으로 조금만 생각하면 쉽게 풀 수 있었다.

일단 뒤에서부터 배열을 읽어나가는 게 핵심이다! 그리고 F[i]의 값을 스택에 push하되, 만약 현재 읽는 배열의 F[i]값이 스택의 top보다 작다면 NGF(i)는 스택의 top이 된다.

만약 스택의 top보다 현재 F[i]가 같거나 크다면, F[i]보다 큰 값이 나올 때까지 스택을 pop한다.

스택이 비어있다면, -1을 출력한다.

 

import sys
n = int(sys.stdin.readline())
arr = list(map(int, sys.stdin.readline().split()))

F = [0]*(1000001)
for i in arr:
    F[i]+=1

stack = []
result = [0]*n

#배열의 뒤에서부터 [F[arr[i]], arr[i]] 스택에 push
for i in range(n-1, -1, -1):
    
    #스택의 top의 F[i]값이 현재 배열의 F[i]값보다 클 때까지 pop
    while(len(stack)!=0 and stack[-1][0]<=F[arr[i]]):
        stack.pop()
        
    #스택이 비어있으면 result에 -1 삽입
    if (len(stack)==0):
        result[i]=-1

    #스택이 비어있지 않은 경우 스택 top의 arr[i]를 result에 삽입    
    else:
        result[i] = stack[-1][1]
    stack.append([F[arr[i]], arr[i]])

#답 출력
for i in result:
    print(i, end=" ")

while문을 돌 때, 스택의 top값을 비교하기 전에 스택이 비어있는지를 먼저 확인하는 조건을 써야 한다.

F의 크기를 선언할 때 처음에 n+1로 선언하는 바보같은 실수를해서 인덱스에러가.. A[i]의 최댓값보다 1만큼 크게 설정하도록 바꿨다.