[17298] 오큰수 | 백준 골드 IV

🔎 문제설명

크기가 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,000,000 이기 때문에 이중 for루프 돌리면 시간초과가 나옵니다.

🧪 자바 코드

import java.io.*;
import java.util.Stack;
import java.util.StringTokenizer;

public class Main {
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));

        Stack<Integer> stack = new Stack<Integer>();

        int N = Integer.parseInt(br.readLine());
        int[] seq = new int[N];


        StringTokenizer st = new StringTokenizer(br.readLine(), " ");

        for(int i = 0; i < N; i++) {
            seq[i] = Integer.parseInt(st.nextToken());
        }


        for(int i = 0; i < N; i++) {
            while(!stack.isEmpty() && seq[stack.peek()] < seq[i]) {
                seq[stack.pop()] = seq[i];
            }
            stack.push(i);
        }

        while(!stack.isEmpty()) {
            seq[stack.pop()] = -1;
        }

        StringBuilder sb = new StringBuilder();
        for(int i = 0; i < N; i++) {
            sb.append(seq[i]).append(' ');
        }

        bw.write(sb.toString());
        bw.flush();
        bw.close();
    }
}

배열을 사용한 코드 (스택보다 빠름)

import java.io.*;
import java.util.StringTokenizer;

public class ObigNum {
    public static void main(String[] args) throws IOException {

        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));

        var N = Integer.parseInt(br.readLine());
        var seq = new int[N];
        var stack = new int[N];	// 스택처럼 쓸 배열
        var top = -1;	// 스택의 맨 위 원소를 가리키는 변수

        StringTokenizer st = new StringTokenizer(br.readLine(), " ");
        for (var i = 0; i < N; i++) {
            seq[i] = Integer.parseInt(st.nextToken());
        }

        for (var i = 0; i < N; i++) {
            while (top != - 1 && seq[stack[top]] < seq[i]) {
                seq[stack[top--]] = seq[i];
            }
            stack[++top] = i;
        }

        for(int i = top; i >= 0; i--) {
            seq[stack[i]] = -1;
        }

        StringBuilder sb = new StringBuilder();
        for (int v : seq) {
            sb.append(v+" ");
        }

        bw.write(sb.toString());
        bw.flush();
        bw.close();
    }
}

🧊 파이썬 코드

import sys


def read(): return sys.stdin.readline().rstrip()


read()  # N
ans = []  # 출력용 리스트
stack = []  # 인덱스 저장 스택
for v in reversed(list(map(int, read().split(' ')))):  # 뒤에서부터 탐색
    while stack and stack[-1] <= v:
        stack.pop()

    ans.append(stack[-1] if stack else -1)  # stack에 존재하지 않으면 가장 큰수이기때문에 -1
    stack.append(v)


print(' '.join(map(str, ans[::-1])))

위에 자바 코드는 앞에서부터 탐색한거라면 파이썬 코드는 뒤에서부터 탐색한 로직입니다.

여러번 코드 바꾸면서 해봤는데 뒤에서 탐색하는 게 시간효율이 안좋게 나왔습니다.. 😂