[백준] 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

예제출력1

5 7 7 -1

---

예제입력2

4
9 5 4 8

예제출력2

-1 8 8 -1

---

문제풀이 

import sys

n = int(sys.stdin.readline())

s = list(map(int, sys.stdin.readline().rstrip().split()))

answer = [-1] * n  # 기본적으로 -1로 초기화된 리스트

먼저 설계를 해보자.

먼저 시간복잡도를 줄이기 위해 input대신에 sys.stdin.readline을 사용한다. 

처음에 나는 for문 안에 for문을 써서 코드를 작성했는데 시간복잡도에 의하여 시간초과라는 결과를 얻었다.

그래서 for문안에 while문으로 해결하기로 했다.

입력받는 횟수 n을 정하고

s라는 리스트에 수열을 입력받는다.

정답으로 제출할 answer리스트를 -1로 채워서 작성한다.

 

stack = []

for i in range(n):

    while stack and s[stack[-1]] < s[i]:

        idx = stack.pop()

        answer[idx] = s[i]

    stack.append(i)

예로 n=4이고 s=[3,5,2,7]이라고 해보자. 그럼 asnwer=[-1,-1,-1,-1]이다.

임의의 빈 스택을 만든다. stack=[], 이 stack은 s의 리스트 인덱스라고 보면 된다.

1. for문 순서대로 i를 stack에 넣는다. 첫 i는 whlie을 안거치고 바로 stack에 append된다.

2. 다음 i 부터 해당 수열보다 s[i]가 클 때를 조건으로 while문을 돌린다.

3. 조건이 맞으면 while문 안에서 stack값을 pop하고 idx에 저장,  answer[idx]=s[i]로 저장한다.

4. while문에 적용되지 않더라도 stack값에 저장해 놨다가, 나중에 조건에 적용되면 그때 pop해서 answer[idx]=s[i]로 저장 가능하다.

5. 여기서 조건이 해당 안되면 stack값은 그대로 남아있게 되고 answer의 요소는 -1로 유지될 것이다. 

 

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

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

map함수를 이용해서 str로 자료형을 바꾼뒤 join함수에 적용해서 한칸씩 띄어서 출력한다. 

정답

import sys

n = int(sys.stdin.readline())
s = list(map(int, sys.stdin.readline().rstrip().split()))
answer = [-1] * n  # 기본적으로 -1로 초기화된 리스트

stack = []
for i in range(n):
    while stack and s[stack[-1]] < s[i]:
        idx = stack.pop()
        answer[idx] = s[i]
    stack.append(i)

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