[Baekjoon] 17298 오큰수(파이썬 python)

이번 문제는 아래 링크에서 풀어볼 수 있습니다.

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().rstrip())
nums = list(map(int, sys.stdin.readline().rstrip().split()))


def mySolution(n, nums):
    stack = []

    for i in range(n):
        compared = -1

        for j in range(i + 1, n):
            if nums[j] > nums[i]:
                compared = nums[j]
                break
            else:
                compared = -1

        stack.append(str(compared))

    return " ".join(stack)


print("mySolution", mySolution(n, nums))

 

모범 답안
파이썬 Ver.

import sys

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


def answer(n, nums):
    nge = [-1] * n
    stack = []

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

    print("answer", *nge)


answer(n, nums)

자바스크립트/NodeJS Ver.

const fs = require('fs');
const filePath = process.platform === 'linux' ? '/dev/stdin' : './input.txt';
let [n, ...input] = fs.readFileSync(filePath).toString().trim().split('\n');

n = parseInt(n);
const data = input[0].split(' ').map((val) => parseInt(val));

const solution2 = (data) => {
    let nge = Array(n).fill(-1);
    const stack = [];

    for (let i = 0; i < n; i++) {
        while (stack.length && data[stack[stack.length - 1]] < data[i])
            nge[stack.pop()] = data[i];
        stack.push(i);
    }

    return nge.map((val) => String(val)).join('\n');
};

console.log('solution2');
console.log(solution2(data));

 

---

문제 풀이

 스택과 인덱스를 활용하여 풀어야 하는 문제다. 나의 경우 이중 for문으로 탐색하는 값보다 오른쪽에 있는 값들을 확인하면서 탐색하는 값보다 큰 값이 나올 경우 스택에 넣고 breack, 없을 경우 -1을 스택에 쌓은 다음 반환하는 방식으로 접근했다. 이 같은 방법은 이중 for문을 돌아야 해서 시간초과가 나왔다.

 스택을 활용하면 이 같은 문제를 해결할 수 있다. 우선 입력 받은 숫자 배열을 for 반복문으로 돈다. 조건문에 부합하지 않을 경우 현재의 인덱스를 스택에 넣는다. 조건에 부합한다면 while 반복문에 진입한다. stack에 값이 있고 스택의 맨 위에 쌓인 인덱스에 해당하는 nums 값과 현재 탐색 중인 nums 값을 비교하여 현재 탐색 중인 값이 큰 경우, 즉 오른쪽에 있는 값이 클 경우 nge리스트에 현재 탐색 중인 값을 삽입하는데 위치는 stack에 쌓인 인덱스를 pop한 인덱스이다. [3, 5, 2, 7]이라는 숫자가 주어졌다면 작업 순서는 다음과 같다.

 

# i = 0
stack이 비어 있으므로 조건문에 부합하지 않는다.
따라서 stack에 현재 인덱스인 0을 append한다.

# i = 1
stack = [0] 이고 nums[stack[-1]] = 3, nums[i] = 5이므로 오큰수의 조건에 부합한다.
stack을 pop한 nge의 0번째 인덱스에 5를 집어넣는다. # nge = [5, -1, -1, -1]
while 반복문의 작업을 마치고 stack에 현재 인덱스 1을 append한다. # stack = [1]

# i = 2
stack = [1] 이지만 nums[stack[-1]] = 5, nums[i] = 2이므로 오큰수의 조건에 부합하지 않는다.
따라서 while 반복문에는 진입하지 않고 stack에 현재 인덱스 2를 append한다. # stack = [1, 2]

# i = 3
stack = [1, 2]이고 nums[stack[-1]] = 2, nums[i] = 7이므로 오큰수의 조건에 부합한다.
stack을 pop한 nge의 2번째 인덱스에 7를 집어넣는다. # nge = [5, -1, 7, -1]
stack = [1]이고 nums[stack[-1]] = 5, nums[i] = 7이므로 오큰수의 조건에 부합한다.
stack을 pop한 nge의 1번째 인덱스에 7를 집어넣는다. # nge = [5, 7, 7, -1]
for 반복문이 끝났고 stack이 비어 더 이상 조건에 부합하지 않으므로 while 반복문에서도 빠져나온다.

 

반복문을 마치고 * 키워드로 리스트 nge의 요소를 분해하여 출력한다.

---

https://aigong.tistory.com/406

[Baekjoon] 17298번 : 오큰수 (Python)