[python] 오큰수

haremeat·2021년 11월 24일
1

Algorithm

목록 보기
15/22
post-thumbnail

백준 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이다.

풀이

금방 해결하겠다 싶었는데 생각보다 푸는 데 오래 걸린 문제였다..!
접근방법은 스택을 하나 만들어서 각 번호의 index값을 넣는다.
만약 해당 번호의 오큰수를 찾으면 stack에서 pop하고 결과값에 넣고
찾지 못하면 그대로 둬서 나중에 스택에 남아있는 index에 해당하는 값들은 전부 -1로 넣어주는 방식으로 풀면 된다.

예를 들어 A = [3, 5, 2, 7]일 때,
일단 0번째 값을 순회하니까 스택에 0을 넣는다.
3의 바로 다음 값인 5가 3보다 크니까
이 경우엔 넣었던 stack = [0]이 pop()되면서 stack = []가 되고 result[0]의 값은 5가 된다.
1번째 값을 순회하니까 stack = [1]이 되고, 5보다 2가 작으므로 이 경우엔 오큰수를 아직 찾지 못했기 때문에 아무것도 하지 않는다.
2번째 값을 순회할 때 stack = [1, 2]가 되고, 2보다 7이 크니까 result[2] = 7이 된다. 오큰수를 찾았기 때문에 stack을 pop()하면서 stack = [1]이 된다.
이 경우 result[1]의 값을 -1로 넣어주면서 끝내면 안 된다. 오큰수는 무조건 인접한 경우에만 구할 수 있는 게 아니기 때문이다. 2의 오큰수는 7이지만 5의 오큰수 역시 7이다.
그래서 오큰수를 찾지 못해 남아있는 값들 역시 체크를 해주면서 마지막까지 스택에 남아있는 값들은 모두 -1로 처리해주면 된다.
당연히 마지막 수인 7은 무조건 -1이 나오기 때문에 굳이 오큰수를 찾을 필요가 없다.

오답 1 (시간초과)

import sys

n = int(sys.stdin.readline())
arr = list(map(int, sys.stdin.readline().split()))
stack = []
result = []

for i in range(n):
    stack.append(i)
    if i < n - 1 and arr[i] < arr[i + 1]:
        stack.pop()
        result.insert(i, arr[i + 1])

        if stack and arr[stack[-1]] < arr[i + 1]:
            result.insert(stack[-1], arr[i + 1])
            stack.pop()

# 스택 값이 남아있다면
if stack:
    for i in stack:
        result.insert(i, -1)

for i in result:
    print(i, end=" ")

첫 번째 오답. 바로 시간초과가 떴다.
for문은 각자 도니까 시간복잡도는 O(n)일텐데 대체 왜 시간초과가 뜨지 ?? 하며 입출력도 바꿔보는 등 온갖 짓을 다 해봤는데 안 풀려서 뭐가 원인인가 싶었더니

result.insert(i, arr[i + 1])

이게 원인이었다. insert()는 배열을 재정렬하기 때문에 시간복잡도 O(n)을 갖는다.
이게 for문 안에 있으니 실제 시간복잡도는 O(n²)이었다...

오답 2 (로직오류)

import sys

n = int(sys.stdin.readline())
arr = list(map(int, sys.stdin.readline().split()))
stack = []
result = [-1] * n

for i in range(n):
    stack.append(i)
    if i < n - 1 and arr[i] < arr[i + 1]:
        stack.pop()
        result[i] = arr[i + 1]

        if stack and arr[stack[-1]] < arr[i + 1]:
            result[stack[-1]] = arr[i + 1]
            stack.pop()

# 스택 값이 남아있다면
if stack:
    for i in stack:
        result[i] = -1

for i in result:
    print(i, end=" ")

그래서 result[i] = arr[i + 1] 방식으로 바꿨더니 시간초과는 안 뜨는데 이번엔 걍 틀렸다고 뜬다. 예시를 이것저것 만들어서 넣어봐도 제대로 나와서 어느 부분이 문제인지 찾느라 오래 걸렸는데 반례가 있었다.

5
[1 1 1 3 4]

위 예시를 넣으면 올바른 답은 3 3 3 4 -1이 떠야하지만 내 코드로는 4 3 3 4 -1이 떴다.
이때 어느 부분이 문제인지 알았다. 그동안은 짧은 배열들만 넣어서 몰랐는데 stack에 남아있는 오큰수를 아직 구하지 못한 값들을 그때그때 전부 돌려줘야 첫번째 값인 1도 오큰수인 3을 구할 수 있다.

제출 코드

import sys

n = int(sys.stdin.readline())
arr = list(map(int, sys.stdin.readline().split()))
stack = []
result = [-1] * n

for i in range(n):
    stack.append(i)
    if i < n - 1 and arr[i] < arr[i + 1]:
        stack.pop()
        result[i] = arr[i + 1]

        while stack and arr[stack[-1]] < arr[i + 1]:
            result[stack[-1]] = arr[i + 1]
            stack.pop()


# 스택 값이 남아있다면
if stack:
    for i in stack:
        result[i] = -1

for i in result:
    print(i, end=" ")

그래서 원래 if문이었던 걸 while로 바꿔줘서 성공!
정말 PS는 꼼꼼해야한다는 걸 깨달은 문제였다...

profile
버그와 함께하는 삶

0개의 댓글