[백준] 17298번 오큰수 (Python)

https://www.acmicpc.net/problem/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)을 공백으로 구분해 출력한다.

예제 입력	예제 출력
4 3 5 2 7	5 7 7 -1
4 9 5 4 8	-1 8 8 -1

---

틀린 아이디어

num 리스트의 뒷 부분부터 접근하여 현재 가리키고 있는 num의 뒷부분에 자신보다 큰 수가 있으면 nge로 설정해주는 것이다.

틀린 코드와 틀린 이유

n = int(input())
nums = list(map(int, input().split()))
# nge를 역순으로 담는 리스트
nge_reverse = [-1]
# pop한 수들을 차례로 담는 리스트
pop_nums = []
while len(nums) > 1:
    # pop_nums 리스트에 pop한 수를 담음
    pop_nums.append(nums.pop())
    # pop_nums리스트의 원소에 순서대로 접근하면서 가장 먼저 나오는 num보다 큰 수를 nge로 설정
    for i in range(-1, -len(pop_nums) - 1, -1):
        # num의 마지막 원소가 pop_num의 원소보다 작다면 nge설정 후 break
        if pop_nums[i] > nums[-1]:
            nge_reverse.append(pop_nums[i])
            break
    # break 없이 끝났다면 num보다 큰 수가 없는 것이므로 nge는 -1
    else :
        nge_reverse.append(-1)
print(*reversed(nge_reverse))

시간초과가 나서 틀린 코드이다. 첫 번째 반복문과 두 번째 반복문으로 인해 시간복잡도가 O(N^2)이다. 

처음에는 다른 코드를 작성했었는데 잘못된 로직이었다. 그 코드를 수정하고 수정하다보니, 굳이 필요 없는 뒤에서부터 접근하는 방식을 사용하게 되었다. 이제는 시간복잡도 계산하는 방법을 조금 잘 알게된 것 같다. 고민하고 고민하다가 모르겠어서 질문게시판에서 힌트를 얻었다. 그 답변글을 보자마자 확 깨달아버렸다.

위 코드는 2중 for 문 (9번째 줄 while + 11번째 줄 for)을 돌면서 총 시간복잡도가 O(n ^ 2) 이 되면서 시간초과가 발생합니다.
 
스택을 사용하면 이미 자신의 오큰수를 발견한 수들을 스택에서 빼주면서
각 원소들을 딱 2 번씩만 스캔해서 시간복잡도를 O(n) 으로 줄일 수 있습니다.
(각 원소들을 스택에 넣을 때 한 번, 스택에서 뺄 때 한 번 총 2번 스캔합니다. 2번 스캔하지만 시간복잡도는 O(n) 이 됩니다.)
스택 사용한 풀이 참고해보세요 (아니면 질문해주세요.)

위의 글을 보고 어떻게 풀어내야 할 지 바로 구상할 수 있었다. 깨닫고 나니까 쉽게 되었던 것 같다. (좋은 답변 남겨주신 kwokws0627님 감사합니다 ㅎㅎ)

풀이 및 코드

n = int(input())
nums = list(map(int, input().split()))
nge_lst = [-1] * n
remains = []
i = 0
while i < n - 1:
    next = nums[i + 1]
    # 어떤 수가 다음 수보다 크거나 같으면 "나머지" 리스트에 추가
    if nums[i] >= next:
        remains.append((i, nums[i]))
    # 다음 수가 더 큰 것이 나오면 아래 코드 실행
    else :
        # 일단 i번째 nge는 next로 설정.
        nge_lst[i] = next
        # 나머지를 하나하나 돌아보면서 나머지 중 next보다 작다면 그 수(나머지)의 nge를 next로 설정
        while remains:
            if remains[-1][1] >= next:
                break
            idx, num = remains.pop()
            nge_lst[idx] = next
    i += 1
print(*nge_lst)

앞에서 차례로 접근하면서 만약 다음 수보다 크다면 나머지 리스트(스택)에 잠시 넣어둔다. 결국 스택에는 내림차순으로 담기게 된다. 그러다가 어떤 수의 다음 수가 더 큰 수가 나타난다면 nge의 조건에 맞는 수들만 스택에서 하나하나 pop해주고, nge를 설정해준다.

n은 최대 100만이고, 제한시간은 1초이므로, 시간복잡도는 O(N)이하여야 한다. while문 안에 또 while문이 있어 시간복잡도가 O(N)보다 크게 보일 수 있지만, 두 번째 while문은 결국 스택에 있는 원소에 접근하는 것이고, remains라는 스택에는 같은 원소가 두 번 이상 들어올 수 없다. 따라서 (첫 번째 while문 한 번 실행시 100만번이 아니라) 이 프로그램을 실행하면서 두 번째 while을 반복하는 횟수는 최대 100만번이다. 즉 두 번째 반복문은 사실상 O(1)이라고 봐도 무방하다.

n의 최댓값이 100만이고, 이 프로그램은 그 원소를 한 번 스택에 넣고 한 번 스택에서 빼는 것이므로 계산 횟수는 최대 200만번이다. 즉, 1초만에 충분히 가능하다.

후기

스택 원소를 pop하고 append하는 코드와 관련된 반복문은 시간복잡도에 크게 영향을 미치지 않는다는 것을 배웠다. 지금까지는 스택이 왜 좋은지, pop을 어떻게 활용해 시간을 줄일 수 있을지에 대한 개념이 부족했던 것 같다. 이 문제를 풀고 스택의 장점을 이용하는 기술을 배웠다.