[Python] 뒤에 있는 큰 수 찾기

문제

링크: https://school.programmers.co.kr/learn/courses/30/lessons/154539

프로그래머스

 

정수 배열 numbers가 매개변수로 주어질 때, 모든 원소에 대한 뒷 큰수들을 차례로 담은 배열을 return 하도록 solution 함수 작성

• 뒷 큰수: 배열의 각 원소들에 대해 자신보다 뒤에 있는 숫자 중에서 자신보다 크면서 가장 가까이 있는 수

# 제한 사항
4 ≤ numbers의 길이 ≤ 1,000,000
1 ≤ numbers[i] ≤ 1,000,000

 

풀이

처음 작성한 코드

시간 초과로 실패: for 문 안에 for 문이 있음. 시간복잡도: O(n^2)

def solution(numbers):
    # 기본값들을 -1 로 채워넣음
    answer = [-1] * len(numbers)

    for i, number in enumerate(numbers):
        for k in range(i, len(numbers)):
            # 큰 수가 나오면 인덱스의 값에 해당값 할당
           	if number < numbers[k]:
                answer[i] = numbers[k]
                break
    
    return answer

 

 

개선한 코드

stack 사용하여 for문 안에 for문 사용 안하도록 개선. 시간 복잡도: O(n)

def solution(numbers):
    answer = [-1] * len(numbers)
    index_stack = []

    for i, number in enumerate(numbers):
    	# 현재 숫자보다 작은 인덱스값을 모두 확인하기 위해 while문 사용
        # if 문을 사용하면 바로 직전의 인덱스값만 확인할 수 있는 문제가 있음
        while index_stack and numbers[index_stack[-1]] < number:
            index = index_stack.pop()
            answer[index] = number
            
        # 자기 자신과 비교를 피하기 위해 while문 밑에서 처리
        index_stack.append(i)

    return answer