뒤에 있는 큰 수 찾기

문제 설명

정수로 이루어진 배열 numbers가 있습니다. 배열 의 각 원소들에 대해 자신보다 뒤에 있는 숫자 중에서 자신보다 크면서 가장 가까이 있는 수를 뒷 큰수라고 합니다.
정수 배열 numbers가 매개변수로 주어질 때, 모든 원소에 대한 뒷 큰수들을 차례로 담은 배열을 return 하도록 solution 함수를 완성해주세요. 단, 뒷 큰수가 존재하지 않는 원소는 -1을 담습니다.

제한 사항

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

입출력 예1

입력

[2, 3, 3, 5]

출력

[3, 5, 5, -1]

입출력 예2

입력

[9, 1, 5, 3, 6, 2]

출력

[-1, 5, 6, 6, -1, -1]

문제 풀이1(시간 초과 방식)

우선 해당 원소의 뒤에 있는 배열중 가장 가깝게 큰 것을 찾는 문제라서, numbers에서 하나의 원소를 고르고 난 뒤, 그 뒤부터 쭉 검사를 했고 가장 먼저 큰 수를 찾으면 answer 배열에 삽입을 했다. 그러다 마지막 원소가 된다면 answer 배열에 -1을 삽입했다. 작은 표본에서는 이게 성립이 되지만, numbers의 길이는 최대 100,000 이었다. 위와 같은 방식을 채택했다면 어마어마한 시간이 걸릴 것 이었다.

코드1(시간 초과 방식)

# 시간 초과
def solution1(numbers):
    answer = []
    count = 0
    lenNumbers = len(numbers)
    for item in numbers:
        count += 1
        if item == max(numbers):
            answer.append(-1)
            continue

        if count == lenNumbers:
            answer.append(-1)
        else:
            flag = False
            for i in range(count, lenNumbers):
                if item < numbers[i]:
                    answer.append(numbers[i])
                    flag = True
                    break

            if flag == False:
                answer.append(-1)

    return answer

print(solution1([2, 3, 3, 5]))
print(solution1([9, 1, 5, 3, 6, 2]))

문제 풀이2

stack을 사용해서 하면 좋을 것 같다는 생각이 들었다. 스택을 사용하면서 얻는 이점은 같은 값을 갖는 원소에 대해서 문제 풀이1과 같은 방법을 하지 않아도 된다. 이게 무슨 뜻이냐면 입출력 예1의 경우 1, 2번 인덱스의 3에 대해서 뒤의 원소들을 추가 분석할 필요가 없다는 뜻이다. 스택에 비교하고자 하는 인덱스를 넣고, 스택이 비어있지 않다면 비교하는 방식을 채택하면 된다. 그러면 스택에 쌓여있는 수들은 모두 뒷 큰수를 만나면 answer 배열을 변경해주기만 하면 되기 때문이다.

코드2

def solution2(numbers):
    stack = []
    answer = [0] * len(numbers)

    for i in range(len(numbers)):
        while stack and numbers[stack[-1]] < numbers[i]:  # 스택이 있는 경우 그리고 그 스택의 가장 위 값이 비교하는 숫자보다 작은 경우
            answer[stack.pop()] = numbers[i]  # 정답 인덱스에 작은 숫자 삽입
        stack.append(i)  # 스택에 인덱스 삽입
    while stack:  # 마지막 원소에 대해서 -1 처리
        answer[stack.pop()] = -1

    return answer

print(solution2([2, 3, 3, 5]))
print(solution2([9, 1, 5, 3, 6, 2]))