[프로그래머스] 뒤에 있는 큰 수 찾기

코딩테스트 연습 - 뒤에 있는 큰 수 찾기

문제 설명

def solution(numbers: list) -> list:
        ...

정수로 이루어진 배열 numbers가 주어진다. 각 원소들에 대해 자신보다 뒤에 있는 숫자 중에서 자신보다 크면서 가장 가까이 있는 수를 뒷 큰수라고 한다. numbers에서 각 원소의 뒷 큰수를 담은 배열을 구하는 문제이다. 만약 뒷 큰수가 존재하지 않으면 -1을 넣는다.

제한사항

• 4 ≤ numbers의 길이 ≤ 1,000,000
  • 1 ≤ numbers[i] ≤ 1,000,000

풀이

복잡하게 풀려고 시도했지만 생각보다 간단한 방법으로 해결할 수 있었다. 접근 방법은 numbers 배열을 순회하며 숫자를 스택에 저장하고 스택에 담긴 숫자들의 뒷 큰수를 갱신하는 식으로 해결했다.

첫 번째 입력예제인 [2, 3, 4, 5] 를 예시로 풀어보면 아래와 같다.

먼저 뒷 큰수를 저장할 answer 배열을 numbers 배열과 동일한 크기로 만들고 -1로 초기화한다.

그리고 numbers 의 인덱스 0부터 순회하며 아래 조건들을 확인하고 스택에 저장한다.

• 현재 인덱스가 가리키고 있는 numbers[i] 의 값과 스택에 저장된 값들을 비교한다.
• 현재의 값이 더 클 경우 이 값이 스택에 저장되있던 값의 뒷 큰수가 된다.

인덱스 순회 과정을 그림으로 나타냈다.

• index 0 : 스택에 현재 인덱스를 추가한다

• index 1
  • numbers[i]가 스택에 있는 값보다 크기 때문에 스택에 저장된 인덱스의 뒷 큰수가 된다. 따라서 answer 배열을 업데이트한다.
  • 현재 인덱스를 스택에 추가한다.

• index 2 : 스택에 현재 인덱스를 추가한다

)

• index 3
  • numbers[i]가 스택의 top보다 크기 때문에 스택에 저장된 인덱스의 뒷 큰수가 된다. 따라서 answer 배열을 업데이트한다.
  • 스택에 남아있던 값이 동일하게 answer 배열을 업데이트한다.
  • 스택에 현재 인덱스를 추가한다.

순회가 끝나면 answer 배열에 뒷 큰수가 갱신되어 있다.

제출 코드

def solution(numbers: list) -> list:
    answer = [-1 for _ in range(len(numbers))]

    s = []
    for i, n in enumerate(numbers):

        while s and numbers[s[-1]] < n:
            answer[s[-1]] = n
            s.pop()

        s.append(i)

    return answer