[BOJ][🟡4][백준#17298] 오큰수

문제 출처

BAEKJOON Online Judge #17298

문제

크기가 N인 수열 A = A1, A2, …, AN이 있다. 수열의 각 원소 Ai에 대해서 오큰수 NGE(i)를 구하려고 한다. Ai의 오큰수는 오른쪽에 있으면서 Ai보다 큰 수 중에서 가장 왼쪽에 있는 수를 의미한다. 그러한 수가 없는 경우에 오큰수는 -1이다. 예를 들어, A = [3, 5, 2, 7]인 경우 NGE(1) = 5, NGE(2) = 7, NGE(3) = 7, NGE(4) = -1이다. A = [9, 5, 4, 8]인 경우에는 NGE(1) = -1, NGE(2) = 8, NGE(3) = 8, NGE(4) = -1이다.

입력

첫째 줄에 수열 A의 크기 N (1 ≤ N ≤ 1,000,000)이 주어진다. 둘째 줄에 수열 A의 원소 A1, A2, …, AN (1 ≤ Ai ≤ 1,000,000)이 주어진다.

출력

총 N개의 수 NGE(1), NGE(2), …, NGE(N)을 공백으로 구분해 출력한다.

예제

예제 1

입력

4
3 5 2 7

출력

5 7 7 -1

예제 2

입력

4
9 5 4 8

출력

-1 8 8 -1

My Sol

import sys
input = sys.stdin.readline
from heapq import heappop, heappush

N = int(input())
arr = list(map(int, input().split()))

Q = []
ret = [-1]*N
for i in range(N):
    heappush(Q, [arr[i], i])
    while Q:
        val, idx = heappop(Q)
        if arr[i] > val:
            ret[idx] = arr[i]
        else:
            heappush(Q, [val, idx])
            break

print(*ret)

heap 사용

1. 주어진 수들을 한 배열에 처리
2. 0부터 N-1까지의 i에 대해 매번 값과 인덱스를 heapQ에 삽입
3. Q에서 현재까지 누적 수 중에 가장 값이 작은 값과 해당 값의 인덱스 추출
4. arr[i]가 val보다 크다면, 어떤 인덱스의 오큰수를 처음으로 발견한 것. ret[idx]에 arr[i]를 부여(오큰수 부여)
5. 이를 반복

결과

맞았습니다!!

모범답안

출처

N = int(input())
nums = list(map(int,input().split(' ')))
out = [-1]*N

stack = [0]
for i in range(1,N):
    while stack and nums[stack[-1]] < nums[i]:
        j = stack.pop()
        out[j] = nums[i]
    stack.append(i)

print(" ".join(map(str,out)))

연산시간이 나의 25% 정도인 풀이를 분석해보려 한다.

이 분은 stack을 사용하셨고, stack에 index만 넣으며 비교하였다. heapq에 원소와 인덱스를 넣고 빼고 할 필요가 없이 stack의 마지막 값만 비교하였으므로 불필요한 과정이나 데이터를 추가로 적재하지 않아 연산시간이 짧았을 것이다.

한 가지 아쉬운 건 입출력에 대해서 *out을 하지 않고 join으로 풀어냈다거나 .split()으로 메서드를 처리하면 .split(“ “)처럼 공백을 default로 한다는 사실을 모르고 계신 것 같다. 그런 것에 비해 알고리즘 자체는 발상이 좋았다.