PS/BOJ

[백준 17298번] 오큰수 파이썬 풀이

76 2024. 1. 19. 00:24
728x90

오큰수 

1 초 512 MB 78543 27776 19581 33.844%

문제

크기가 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

 

그동안 그래프 탐색 문제만 계속 풀어서 자료구조 탭에서 골라서 풀어본 문제이다.

처음 문제를 보고 생각했던 풀이는 2중 반복문을 이용해서 배열의 처음부터 끝까지 탐색을 진행하며 인덱스의 값보다 큰 수 중 가장 왼쪽에 있는 수를 단순히 일일히 계산하였으나, 이렇게 문제를 풀 경우 O(n^2)의 시간이 걸리기 때문에 시간초과 판정을 받을수 밖에 없었다.

 

아래는 첫번째로 시도했던 방법이다.

n = int(input())
data = list(map(int, input().split()))

for i in range(n):
    num = data[i]
    nge = []
    for j in range(i+1, n):
        if data[j] > data[i]:
            nge.append(data[j])
            break
    if nge:
        print(nge[0], end=' ')
    else:
        print(-1, end=' ')
 
 

풀이 방법은 다음과 같다.

 

  1. 스택과 NGE배열을 준비한다. 스택에는 배열의 0번째 값을 미리 삽입해두고, NGE배열의 값은 모두 -1로 초기화한다.
  2. 이어서 2번째 인덱스부터 배열의 값을 차례로 스택의 가장 위의 값과 대소비교를 진행한다.
  3. 배열의 값이 스택의 값보다 크다면, 해당 값이 스택의 값의 오큰수가 된다.(오큰수는 NGE배열에 기록한다.)
  4. 대소비교를 스택의 데이터보다 배열의 값이 작아질때까지 반복한다.
  5. 배열의 값을 스택에 다시 삽입한다.

 

그림을 통해 설명하면 아래와 같다.

왼쪽에서부터 차례로 데이터 배열, 스택, 오큰수 배열이다.

0번째 인덱스 9는 이미 스택에 넣어놓은 상태로 알고리즘을 시작한다.

5는 9보다 작은수 이므로 스택에 삽입한다.

이어서 4와 5를 비교한다. 4는 5보다 작은수이므로 스택에 삽입한다.

이어서 8과 4를 비교한다. 8은 4보다 크므로, 4의 오큰수는 8이 된다. 4의 인덱스는 data배열에서 3에 해당하므로 NGE[3]의 값을 8로 수정하고 4를 스택에서 팝한다.

이어서 8과 5를 비교한다. 8은 5보다 크므로, 5의 오큰수는 8이 된다. 5의 인덱스는 data배열에서 2에 해당하므로 NGE[2]의 값을 8로 수정하고 5를 스택에서 팝한다.

이어서 8과 9를 비교한다. 8은 9보다 작으므로 스택에 8을 삽입한다.

더이상 비교할 대상이 없으므로, 알고리즘을 종료한다. NGE를 출력하게되면 9의 오큰수는 없으므로 -1, 5의 오큰수는 8, 4의 오큰수는 8, 8의 오큰수는 없으므로 -1을 출력하게 된다.

 

위 과정을 코드로 옮기면 다음과 같다.

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

n = int(input())
data = list(map(int, input().split()))

stack = [0]  # 알고리즘에 사용할 스택, 첫번째 값은 미리 스택에 넣어둔다.

nge = [-1]*n

# 인덱스를 스택에 삽입할것임
for i in range(1, n):
    # 스택에 데이터가 존재하면서, 스택의 가장 위에있는값이 현재 삽입하려는 데이터보다 작은값이라면
    # 스택의 가장위의 값은 stack[-1]로 표현가능
    while stack and data[stack[-1]] < data[i]:
        # 오큰수 기록
        nge[stack.pop()] = data[i]  # 이후에 삽입하려는 데이터가 더 크다면 해당수는 오큰수가 됨
    stack.append(i)  # 이후 i번째 인덱스 삽입

# 배열 출력시 *을 이용하면 공백을 기준으로 출력됨
print(*nge)