[레벨 2] 뒤에 있는 큰 수 찾기

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

프로그래머스

주어진 숫자 리스트에서 i 번째 원소 뒤에 있는 원소 중 i보다 큰 가장 가까운 원소를 찾아 리스트를 만드는 문제이다.

예를 들어, numbers = [2,3,3,5]이면, 아래와 같은 과정을 통해 정답을 계산한다.

1. index = 0: 2 뒤에서 2보다 큰 가장 가까운 원소는 index = 1인 3

2. index = 1: 3 뒤에서 3보다 큰 가장 가까운 원소는 index = 3인 5

3. index = 2: 3 뒤에서 3보다 큰 가장 가까운 원소는 index = 3인 5

4. index = 3: 5 뒤에서 5보다 큰 가장 가까운 원소는 없으므로 -1

그러므로 정답은 [3,5,5,-1]이 된다.

 

먼저, 0번째부터 -2번째까지의 원소를 훑으며, 해당 원소 뒤의 원소들을 훑으며 큰 원소가 나오면 break하는 코드를 작성했다.

즉, for loop을 2번 돌리는 O(n^2) 연산량의 코드를 작성했다.

def solution(numbers):
    answer = []
    for i in range(len(numbers)-1):
        for j in range(i+1,len(numbers)):
            if numbers[i] < numbers[j]:
                answer.append(numbers[j])
                break
        if len(answer) != i+1:
            answer.append(-1)
    answer.append(-1)
    return answer

의외로 대부분의 테스트 케이스에서 통과해서 놀랐으나, 마지막 4개의 테스트는 시간 초과로 실패했다.

약간의 개선으로 통과할 수 있을까 싶어서, 해당 원소 뒤의 원소 중 중복 원소를 제거하는 방식으로 코드를 약간 수정했다.™

# i 번째 이후 원소 중복 제거 및 순서 유지 후 탐색
def solution(numbers):
    answer = []
    for i in range(len(numbers)-1):
        for j in list(dict.fromkeys(numbers[i+1:])):
            if numbers[i] < j:
                answer.append(j)
                break
        if len(answer) != i+1:
            answer.append(-1)
    answer.append(-1)
    return answer

하지만 결과는 더 안좋아졌다. 딕셔너리를 통해 중복은 없애며 순서는 유지했는데, 연산량이 더 늘어났나보다.

역시 O(n^2)의 연산량으로는 문제를 풀 수가 없었다.

그래서 stack 자료구조를 이용하여 풀었더니 한 번에 통과할 수 있었다.

def solution(numbers):
    stack = []
    answer = [-1 for _ in range(len(numbers))]
    for i,j in enumerate(numbers):
        if i == 0:
            stack.append((i,j))
        else:
            if stack[-1][1] > j:
                stack.append((i,j))
            else:
                while stack != [] and stack[-1][1] < j:
                    k,l = stack.pop()
                    answer[k] = j
                stack.append((i,j))
    return answer

사실 stack을 떠올리질 못해서 질문하기를 통해 stack을 사용하면 된다는 힌트를 얻고 풀었다.

나에게는 앞에 풀었던 카카오 블라인드 테스트 문제보다 어려운 문제였다...