[Java+Python]17298번, 오큰수

목차

 

문제

 

크기가 N인 수열 $A = A_1, A_2,..., A_N$이 있다. 

수열의 각 원소 $A_i$에 대해서 오큰수 NGE(i)를 구하려고 한다. 

$A_i$의 오큰수는 오른쪽에 있으면서 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의 원소 $A_1, A_2,..., A_N (1 ≤ A_i ≤ 1,000,000)$이 주어진다.

 

출력

 

총 N개의 수 NGE(1), NGE(2),..., NGE(N)을 공백으로 구분해 출력한다.

 

풀이

 

그야말로 스택을 위해 만들어진 문제 그 자체이다.

 

자신보다 오른쪽에 있으면서 자신보다 큰 수 중 가장 왼쪽에 있는 수라는 것을 판별하고 꺼내려면 스택이 제격이기 때문이다.

 

다만 이를 위해선 스택에 숫자 자체가 아닌 인덱스를 저장해야 하는데, 해당 수가 수열에서 자리한 위치를 알아야 하기 때문이다.

 

무슨 말인지는 코드를 보면 명확해진다.

		int n = Integer.parseInt(br.readLine());
		int[] arr = new int[n];
		int[] result = new int[n];
		Arrays.fill(result, -1);

먼저 입력 값의 첫 번째 줄을 n으로 받아 arr과 result를 선언하고, result를 -1로 채운다.

 

이는 오른쪽에 큰 값이 없을 경우의 기본값이 -1이기 때문이며, 이렇게 선언하면 뒤에 굳이 -1로 다시 갱신할 필요가 없다.

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

두 번째 줄의 입력값으로 arr을 채운다.

		ArrayDeque<Integer> stack = new ArrayDeque<>();
		stack.push(0);

Stack 자료구조 대신에 ArrayDeque 자료구조를 선택하고 이름을 stack이라 붙인다.

 

그리고 스택의 가장 처음 값으로 0을 집어넣어 1번 인덱스부터 순회할 수 있도록 한다.

 

그 이유는 아래의 글에 나와 있다.

[면접 준비 - Java]Data Structures in a Nutshell

[면접 준비 - Java]Data Structures in a Nutshell

		for (int i = 1; i < n; i++) {
			while (!stack.isEmpty() && arr[stack.peek()] < arr[i]) {
				result[stack.pop()] = arr[i];
			}
			stack.push(i);
		}

계속해서 1번 인덱스부터 n-1까지 반복하며 다음 작업을 수행한다.

 

참고로 n-1까지만 반복하는 이유는 가장 오른쪽 수는 비교할 대상이 없어 무조건 -1이 되기 때문이다.

 

1. 스택이 비어있지 않으면서 스택의 맨 위에 있는 원소를 인덱스로 하는 arr의 값이 arr[i]보다 작은 경우,
   문제의 조건을 만족하는 '오큰수'이므로 result배열의 해당 위치에 현재 요소를 저장한다.
2. 이를 스택이 비거나, 스택의 맨 위 원소가 현재 원소보다 크거나 같을 때까지 반복한다.
3. 현재 원소의 인덱스를 스택에 푸시한다.

		StringBuilder sb = new StringBuilder();
		for (int i : result) {
			sb.append(i).append(' ');
		}

		System.out.println(sb);

마지막으로 스트링 빌더를 사용해 result를 모두 출력한다.

 

위의 로직을 보면 알겠지만, 자신의 오른쪽에 더 큰 수가 존재하지 않는다면 기본값인 -1로 남을 것이다.

 

Java

 

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayDeque;
import java.util.Arrays;
import java.util.StringTokenizer;

public class Main {

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

		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

		int n = Integer.parseInt(br.readLine());
		int[] arr = new int[n];
		int[] result = new int[n];
		Arrays.fill(result, -1);

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

		ArrayDeque<Integer> stack = new ArrayDeque<>();
		stack.push(0);

		for (int i = 1; i < n; i++) {
			while (!stack.isEmpty() && arr[stack.peek()] < arr[i]) {
				result[stack.pop()] = arr[i];
			}
			stack.push(i);
		}

		StringBuilder sb = new StringBuilder();
		for (int i : result) {
			sb.append(i).append(' ');
		}

		System.out.println(sb);
	}
}

 

Python

 

import sys

n = int(sys.stdin.readline().rstrip())
lst = list(map(int, sys.stdin.readline().rstrip().split(" ")))
result = [-1] * n

stack = []
stack.append(0)

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

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

새삼 파이썬은 정말, 정말 코드가 간결하고 가독성이 좋다.

 

Performance