파이썬 코딩테스트 강좌, 오큰수 구하기

안녕하세요! 오늘은 취업 준비생들에게 매우 유용한 알고리즘 문제인 ‘오큰수 구하기’를 소개하고 자세히 설명하는 시간을 가지겠습니다. 이 문제는 스택(Stack)의 개념을 이해하고 활용하는 데 큰 도움이 될 것입니다. 그럼 시작해볼까요?

문제 설명

문제 설명: 주어진 배열에서 각 원소에 대해 “오큰수”를 구하는 문제입니다. “오큰수”는 해당 원소의 오른쪽에 있는 원소 중에서 가장 첫 번째로 크기가 큰 원소를 의미합니다. 만약 그러한 원소가 존재하지 않으면 -1로 표시합니다.

입력: 첫 번째 줄에 자연수 N(1 ≤ N ≤ 1,000,000)이 주어지고, 두 번째 줄에는 N개의 자연수 A₁, A₂, …, A_N (1 ≤ A_i ≤ 1,000,000)이 주어집니다.

출력: 각 원소에 대한 오큰수를 한 줄에 하나씩 출력합니다.

예제

입력:

5
2 1 3 4 5

출력:

문제 해결 전략

문제를 해결하기 위한 가장 효과적인 방법 중 하나는 스택을 이용하는 것입니다. 스택은 LIFO(Last In, First Out) 구조의 자료구조로, 문제를 해결하는 데 유용하게 사용될 수 있습니다.

스택을 이용한 접근 방법

비어 있는 스택을 하나 생성합니다.
주어진 배열을 오른쪽부터 왼쪽으로 탐색하면서 각 원소에 대해 다음을 수행합니다:

현재 원소와 스택의 최상단 원소를 비교합니다.
스택의 최상단 원소가 현재 원소보다 크면, 그 원소가 오큰수입니다. 현재 원소에 오큰수를 저장합니다.
스택의 최상단 원소가 현재 원소보다 작거나 스택이 비어있으면 스택에서 원소를 pop합니다.
현재 원소를 스택에 추가합니다.

모든 원소를 탐색한 후, 오큰수가 저장된 결과 배열을 출력합니다.

코드 구현

이제 위의 알고리즘을 파이썬 코드로 구현해보겠습니다.


def find_next_greater_element(arr):
    n = len(arr)
    result = [-1] * n  # 오큰수를 저장할 결과 배열
    stack = []  # 스택 생성

    # 배열을 오른쪽부터 왼쪽으로 탐색
    for i in range(n - 1, -1, -1):
        # 스택의 최상단 원소와 비교
        while stack and stack[-1] <= arr[i]:
            stack.pop()  # 스택에서 pop
        
        # 스택이 비어있지 않다면 오큰수 기록
        if stack:
            result[i] = stack[-1]
        
        # 현재 원소를 스택에 추가
        stack.append(arr[i])

    return result

# 입력 처리
n = int(input())
arr = list(map(int, input().split()))

# 오큰수 구하기
result = find_next_greater_element(arr)

# 결과 출력
for r in result:
    print(r)

코드 설명

위 코드는 주어진 배열에서 각 원소의 오큰수를 찾는 함수입니다. 함수는 다음과 같은 과정을 따릅니다:

입력으로 받은 배열의 길이를 구하고, 결과를 저장할 배열을 초기화합니다.
주어진 배열을 오른쪽부터 왼쪽으로 순회합니다.
현재 원소와 비교하기 위해 스택의 최상단 원소를 확인합니다.
스택의 최상단 원소가 현재 원소보다 작거나 같으면 스택에서 pop합니다. 이렇게 해서 스택에는 항상 현재 원소보다 큰 값만 남게 됩니다.
스택이 비어 있지 않다면, 스택의 최상단 원소가 현재 원소의 오큰수가 됩니다. 이것을 결과 배열에 기록합니다.
현재 원소를 스택에 추가합니다.

복잡도 분석

시간 복잡도는 O(N)입니다. 각 원소는 스택에 한 번 추가되고 한 번 제거되므로, 전체적으로 N번의 연산이 필요합니다. 공간 복잡도는 O(N)으로 결과 배열과 스택을 저장하는 데 필요한 공간을 고려해야 합니다.

결과 예시

이제 코드의 일부 테스트 케이스를 확인해보겠습니다. 위에서 제시한 예제를 사용하여 오큰수를 구해보겠습니다.


# 입력 예제
# 5
# 2 1 3 4 5

print(find_next_greater_element([2, 1, 3, 4, 5]))  # 출력: [3, 3, 4, 5, -1]

결론

이번 시간에는 ‘오큰수 구하기’ 문제를 해결하는 과정을 살펴보았습니다. 스택을 사용하여 효율적으로 문제를 해결하고, 실제 코드를 통해 그 원리를 적용해 보았습니다. 이러한 방식의 문제 해결은 실제 코딩 테스트와 알고리즘 면접에서 매우 유용하니 꼭 연습해 보시기 바랍니다.

다음 시간에는 또 다른 알고리즘 문제를 가지고 찾아오겠습니다. 감사합니다!