백준-BOJ 17298번 오큰수 파이썬 풀이

Python

[Python] 백준 알고리즘 온라인 저지 17298 오큰수

출처

https://www.acmicpc.net/problem/17298

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

---

문제 설명

주어진 n개의 수들로 이루어진 배열에 각 수들의 자신의 인덱스보다 오른쪽에 있으면서 큰 수를 찾아내야 하는 문제로 시간 초과를 주의해야 하는 문제이다. 처음에는 각 수마다 남아있는 배열을 for문으로 선회하려고 하였지만 시간복잡도를 O(n)으로 만들어야 했고 이 과정에서 스택의 개념을 사용했다.

 

풀이

import sys
input = sys.stdin.readline

n = int(input())
arr = list(map(int, input().split()))
ans = [-1] * n
stack = [0]

for i in range(1, n):
    while stack and arr[stack[-1]] < arr[i]:
        ans[stack.pop()] = arr[i]
    stack.append(i)

print(*ans)

 

풀이 설명

이 풀이에서 핵심적인 내용은 스택에 수들을 직접 넣는 것이 아니라 인덱스를 넣어 그 값을 활용하여야 한다는 점과 앞서 문제 설명에서와 같이 배열을 한 번만 선회하면서 크기를 비교해야 하기에 스택을 사용해야 한다는 점이다.

이 개념에 대해 살펴보면 첫 번째 인덱스인 0을 stack 넣어놓고 시작하여 다음 수부터 배열을 선회하는 for문을 통해 stack 안에 수가 있고 stack의 마지막 수 (인덱스)를 arr와 i번째 수와 비교한다.

 [3, 5, 2, 7]의 경우

• 1. stack = [0], i = 1
  
  • 이 경우 5는 3의 오큰수이기 때문에 바로 stack에서 마지막수 0을 빼고 ans[0]을 수정해준다. 그리고 1을 다시 스택에 넣는다.
• 2. stack = [1], i = 2
  • arr[stack[-1]]이 2보다 크기 때문에 오큰수가 아니므로 다음 인덱스를 스택에 집어넣는다.

• 3. stack = [1, 2], i = 3
  • 이제 스택의 마지막 수 arr[2]가 3번째 인덱스 수보다 작기 때문에 ans를 arr[3]으로 수정하고 스택에서 뺀다.
  • 아직 스택에는 1이 남아있고 arr[3]과 arr[1]을 다시 비교한다 5보다 7이 크기 때문에 오큰수로 ans[1]도 수정한다.
  • 마지막 수는 그냥 넘어간다.

마지막 수와 오큰수가 없는 경우에는 -1을 출력해야 하기에 ans라는 배열을 -1로 미리 설정해놓은 것이다.

 

풀이 결과