[파이썬] 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)을 공백으로 구분해 출력한다.

예제 입력 1 복사

4
3 5 2 7

예제 출력 1 복사

5 7 7 -1

예제 입력 2 복사

4
9 5 4 8

예제 출력 2 복사

-1 8 8 -1

 

생각

제일 오른쪽에 있는 큰수를 찾으라고 했으니까 인덱스비교를 사용하는데 이중 for문을 사용하면 되겠구나!

라고 생각

import sys
n=int(sys.stdin.readline())
a=list(map(int,sys.stdin.readline().split()))
for j in range(n):
    base=a[j]
    if base==max(a) or j==n-1:
        print(-1, end=' ')
        continue
    for i in range(j+1,n+1):
        if i==n:
            print(-1, end=' ')
            break
        elif base<a[i]:
            print(a[i],end=' ')
            break
            
            
            
import sys
n=int(sys.stdin.readline())
a=list(map(int,sys.stdin.readline().split()))
for j in range(n):
    base=a[j]
    for i in range(j+1,n+1):
        if i==n:
            print(-1, end=' ')
        elif base<a[i]:
            print(a[i],end=' ')
            break

위와 같은 코드로 돌렸지만 시간초과가 났다. 

아무리 생각을 해도 되지 않아 검색을 통해 알아 보았다.

 

이중 for문을 쓸 경우에 num의 최댓값이 백만으로 백만 제곱에 해당하는 시간이 걸리게 되어 시간초과가 발생한다는 것이었다.

그래서 stack을 사용하자라는 것으로 접근하였다

왜냐하면 stack을 사용하게 되면 stack이라고 지칭하는 리스트에 인덱스에 관한 것을 넣게된다.

인덱스에 관한 것을 넣는데 

이때의 과정은

stack이라는 변수이름을 가진 stack 안에 숫자가 있고

stack에 저장된 인덱스와 for문으로 증가되고 있는 인덱스에 있는 값들을 비교해서 증가되고 있는 인덱스에 속한 값이 더 크다.

라는 위의 2조건이 맞다면 NGE라는 곳에 저장된 -1에 해당하는 값은 바뀌게 되는 것이다.

만약 틀리다면 stack에 for문으로 증가되고 있는 인덱스만 넣게 되는 꼴이 된다.

 

즉 그냥 반목문으로 이중 for문으로 작성하게 되면 O(N^2)의 시간복잡도를 가질 수 있는 것인데 스택을 사용해 O(N)에 가깝게 만들고자 스택을 사용한다.

시간을 줄이게 되는 부분은 스택에 인덱스를 집어넣고 pop을 하면서 오큰수를 업데이트 하는 것이다. 그러니까 이중 for문을 돌릴 수도 있는 것을 특정 상황인 스택에 인덱스가 집어넣어져있는 상태에만 진행시켜 시간이 줄어들게 된다.

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

NGE = [-1 for _ in range(n)]
stack = []
for i in range(n) :
    while(len(stack) != 0 and a[stack[-1]] < a[i]) :
        NGE[stack[-1]] = a[i]
        stack.pop()
    stack.append(i)
for i in NGE:
    print(i,end=' ')

 

입력

4
3 5 2 7

i	0	1	2	3	x
NGE	-1  -1  -1  -1	5  -1  -1  -1	5  -1  -1  -1	5  -1  7  -1	5  7  7  -1
len(stack)	0	1	1	2	x
stack	0	1	1  2	1	x