[Java+Python]17299번, 오등큰수

목차

 

문제

 

크기가 N인 수열 $A = A_1, A_2,..., A_N$이 있다. 수열의 각 원소 $A_i$에 대해서 오등큰수 NGF(i)를 구하려고 한다.

$A_i$가 수열 A에서 등장한 횟수를 $F(A_i)$라고 했을 때, $A_i$의 오등큰수는 오른쪽에 있으면서
수열 A에서 등장한 횟수가 $F(A_i)$보다 큰 수 중에서 가장 왼쪽에 있는 수를 의미한다.
그러한 수가 없는 경우에 오등큰수는 -1이다.

예를 들어, A = [1, 1, 2, 3, 4, 2, 1]인 경우 F(1) = 3, F(2) = 2, F(3) = 1, F(4) = 1이다. 
$A_1$의 오른쪽에 있으면서 등장한 횟수가 3보다 큰 수는 없기 때문에, NGF(1) = -1이다. 
$A_3$의 경우에는 $A_7$이 오른쪽에 있으면서 F(A3=2) < F(A7=1) 이기 때문에, NGF(3) = 1이다.

NGF(4) = 2, NGF(5) = 2, NGF(6) = 1이다.

 

입력

 

첫째 줄에 수열 A의 크기 N (1 ≤ N ≤ 1,000,000)이 주어진다. 

둘째에 수열 A의 원소 $A_1, A_2,..., A_N (1 ≤ A_i ≤ 1,000,000)$이 주어진다.

 

출력

 

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

 

풀이

 

지난 문제와 비슷하지만 이번엔 숫자의 크기가 아니라 빈도수의 크기를 기준으로 하는 문제다.

 

크게 고민할 것 없이 지난번 문제처럼 풀면 되는데, 요약하자면

 

배열의 각 요소를 돌며 현재 요소의 빈도수가 스택의 top보다 높다면 해당 top 요소를 인덱스로 하는

 

result 배열의 값을 현재 요소로 설정하고 스택에서 top을 pop 한다.

 

이 과정을 스택이 비거나(-1), 처음으로 요소의 빈도가 top보다 높을 때까지 반복한다.

 

코드를 토막내서 읽으며 조금 더 설명을 해보겠다.

public class Prob17299 {
	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[] freq = new int[1000001];
		int[] result = new int[n];
		Arrays.fill(result, -1);

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

입력값으로 주어진 n을 이용해 arr과 result를 초기화하고 빈도수를 저장할 freq을 문제의 조건에 맞게 초기화한다.

 

이어서 result 배열을 -1로 채우는데, 이는 오등큰수를 찾지 못했을 경우를 위한 초기값이다.

 

그리고 늘 그렇듯 어레이덱을 이용해 스택을 선언한다.

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

이어서 다음 입력값을 받아 arr 배열에 채워 넣으며 해당 값의 빈도수를 계산해 freq에 넣어준다.

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

풀이의 핵심이 되는 코드이다. 먼저 스택에 0을 넣고, i를 1부터 시작해서 반복문을 돈다.

 

반복문 안에선 먼저 첫 번째 원소를 스택에 넣고(0보다는 무조건 크기 때문에),

 

다음 바퀴부터 해당 인덱스의 빈도수가 스택의 가장 위에 있는 요소의 빈도수보다 크다면

 

스택에서 pop을 하며 동시에 result[stack.pop()]을 현재 인덱스의 값으로 갱신한다.

 

이를 스택이 비거나 현재 인덱스의 빈도수가 가장 위에 있는 요소의 빈도수보다 작아질 때까지 반복한.

 

이를 반복하면 각 요소에 대해 자신보다 빈도수가 높은 첫 요소를 찾게 된다.

 

이 과정이 끝나면 현재 인덱스를 스택에 push 한다.

 

마지막까지 스택에 남아있는 인덱스는 오등큰수를 찾지 못한 수 이므로, -1로 초기화한 그대로의 값을 가지게 된다.

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

		System.out.println(sb);
	}
}

마지막으로 result 배열을 순회하며 문자열을 생성, 출력한다.

 

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 Prob17299 {

	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[] freq = new int[1000001];
		int[] result = new int[n];
		Arrays.fill(result, -1);

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

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

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

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

		System.out.println(sb);
	}
}

 

Python

 

import sys

n = int(sys.stdin.readline().rstrip())
lst = list(map(int, sys.stdin.readline().rstrip().split(" ")))
freq = [0] * 1000001
for num in lst:
    freq[num] += 1

result = [-1] * n

stack = []
stack.append(0)

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

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

길이가 말 그대로 자바의 절반 밖에 되지 않는다.

 

Performance

 

하지만 메모리와 시간은 역시 많이 든다. 반복문은 어쩔 수 없다.