[스택] 코딩테스트 문제 TIL (오큰수) - 백준 17298번

말하는 감자·2024년 12월 17일

algorithm baekjoon data structure python stack 백준 스택 알고리즘 자료구조 코딩테스트 코딩테스트 준비 파이썬

코딩테스트 유형별 풀이

목록 보기

34/67

최초로 백준 골드 문제를 풀게 되었습니다. 그런데 해당 문제는 LeetCode 739문제와 매우 유사합니다. 해당 문제 정보는 다음 링크에서 확인할 수 있습니다. 그럼 살펴보도록 하겠습니다.

1. 오늘의 학습 키워드

스택

2. 문제: 오큰수

단계: 🥇골드 4단계
주제: 자료 구조, 스택
출처: https://www.acmicpc.net/problem/17298

문제

크기가 N인 수열 A = $A_1, A2_, ..., A_N$ 이 있다. 수열의 각 원소 $A_i$ 에 대해서 오큰수 NGE(i)를 구하려고 한다. $A_i$ 의 오큰수는 오른쪽에 있으면서 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의 원소 $A_1, A_2, ..., A_N$ (1 ≤ $A_i$ ≤ 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

3. 문제 풀이

Step1. 문제 이해하기

크기가 N인 수열 A = $A_1, A2_, ..., A_N$ 이 있습니다. 수열의 각 원소 $A_i$ 에 대해서 오큰수 NGE(i)를 구하려고 합니다. $A_i$ 의 오큰수는 오른쪽에 있으면서 Ai보다 큰 수 중에서 가장 왼쪽에 있는 수를 의미합니다. 그러한 수가 없는 경우에 오큰수는 -1이라고 합니다.

즉 현재 인덱스의 값보다 큰 인덱스 중 가장 가까이(작은) 인덱스에 해당하는 값이 현재 인덱스의 오큰수입니다.

뭔가 감이 오지 않나요?

현재 위치 이후에 값들중 현재 위치 값보다 큰 값에만 반응을 합니다. 과거는 의미가 없는 것이며 현재 위치로부터 미래의 값들과 비교하는것입니다. 바로 이것이 LIFO원리를 이용하는것입니다!

현재 위치 값이 수열 A를 제일 최근에 순회한 값일테니까요!

더 자세히 분석하기 전 입출력 조건을 살펴보도록 하겠습니다.

Input:
- 첫째 줄에 수열 A의 크기 N이 주어집니다.
- N은 1이상 $10^6$ 이하입니다. 따라서 $O(n^2)$ 보다 효율적인 시간 복잡도로 코드를 구현해야 합니다.
- 스택을 사용한다면 수열 A만 순회하면 되기 때문에 충분히 가능할 것으로 보입니다.
- 둘째 줄에 수열 A의 원소 $A_1, A2_, ..., A_N$ 가 주어지며 각 수열의 원소 값은 1이상 $10^6$ 이하입니다.
Output:
- 총 N개의 수 NGE(1), NGE(2), ..., NGE(N)을 공백으로 구분해 출력해야 합니다.

Step2. 문제 분석하기

사실, Step1에서 거의 다 분석이 완료되었습니다.

스택 자료구조를 활용해서 구현하면 되고, 이것은 수열 A의 크기만큼의 시간 복잡도로 구현될 것입니다. 따라서 N 크기에도 충분히 여유가 있는것을 알 수 있습니다.

그렇다면 스택 자료구조로 구체적으로 어떻게 구현하면 될 지 살펴 보도록 하겠습니다.

구체적인 접근법:

스택 활용:
- 인덱스를 스택에 저장합니다.
- 스택의 top에 있는 인덱스의 값과 현재 순회 중인 값을 비교합니다.
- 만약 현재 값이 스택의 top에 해당하는 값보다 크다면 오큰수 조건을 만족한 것이므로, 스택의 top을 pop하고 해당 위치의 오큰수를 업데이트합니다.
반복 수행:
- 스택에서 pop된 인덱스들은 그 값의 오큰수를 이미 찾았으므로, 다음 값을 비교할 필요가 없습니다.
- 현재 인덱스를 스택에 push합니다.
결과 처리:
- 모든 원소를 순회한 후에도 스택에 남아있는 인덱스들은 오큰수를 찾지 못한 경우이므로 1로 설정됩니다.

이 방법의 핵심은 모든 원소를 한 번씩만 검사하고, 스택에서 pop 연산도 $O(1)$ 에 수행되므로 전체 시간 복잡도가 $O(N)$ 입니다.

Step3. 코드 설계

입력값 초기화:
- N과 수열 A를 입력받습니다.
- NGE 배열을 -1로 초기화합니다. (모든 값이 오큰수를 찾지 못한 상태로 시작)
- 빈 스택을 초기화합니다.
수열 순회:
- A의 각 원소를 순회하면서 스택의 top과 비교합니다.
- 조건에 맞으면 스택에서 pop하고 오큰수를 기록합니다.
- 현재 인덱스를 스택에 push합니다.
결과 출력:
- 결과를 출력 형식에 맞게 변환하여 반환합니다.

Step4. 코드 구현

import sys

N = int(sys.stdin.readline().strip())  # 수열의 크기 입력
A = list(map(int, sys.stdin.readline().split()))  # 수열 입력

def sol(N, A):
    NGE = [-1] * N  # 결과 배열을 -1로 초기화
    stack = []  # 인덱스를 저장할 스택 초기화
    
    for i in range(N):
        # 현재 값 A[i]가 스택의 top에 있는 인덱스의 값보다 크면 오큰수를 찾은 것
        while stack and A[stack[-1]] < A[i]:
            NGE[stack.pop()] = A[i]
        stack.append(i)  # 현재 인덱스를 스택에 저장
    
    # 결과를 공백으로 구분된 문자열로 변환하여 반환
    return ' '.join(map(str, NGE))

# 결과 출력
print(sol(N, A))

코드 설명:
- 입력 처리:
  - N과 수열 A를 입력받습니다.
- 결과 초기화:
  - NGE 배열은 1로 초기화되며, 이는 모든 값이 오큰수를 아직 찾지 못했다는 의미입니다.
- 스택 활용:
  - 인덱스를 스택에 저장합니다.
  - 현재 값 A[i]가 스택의 top에 있는 값보다 크다면, 오큰수 조건을 만족한 것이므로 NGE 배열을 업데이트하고 스택을 pop합니다.
  - 조건을 만족하지 않으면 현재 인덱스를 스택에 push합니다.
- 결과 반환:
  - 모든 원소를 순회한 후 NGE 배열의 결과를 공백으로 구분된 문자열로 반환합니다.
시간 복잡도: $O(N)$
결과:

4. 마무리

이 문제는 스택을 활용해 효율적으로 오큰수를 찾는 문제로, 각 원소를 한 번씩만 탐색하므로 O(N)의 시간 복잡도로 해결할 수 있습니다. 스택을 활용하면 과거의 데이터를 저장하면서 현재 값과 비교하는 LIFO구조를 잘 이용할 수 있다는 점이 핵심입니다. 또한, 결과를 저장하는 배열과 스택을 활용해 코드가 직관적이고 간결하게 구성됩니다. 이 문제를 통해 스택을 활용한 탐색 최적화 방법을 다시 한 번 익힐 수 있었습니다.

전체 코드는 다음 링크에서 확인할 수 있습니다.

읽어주셔서 감사합니다!