[알고리즘_코딩 테스트_JAVA] 오큰수 (JAVA, PYTHON)

본문은 알고리즘 코딩테스트 자바편을 정리한 내용입니다 :)

 

✔ 문제 012 오큰수 구하기

⬜ 핵심 아이디어

반복문으로 오큰수를 찾으면 제한시간을 초과한다. 한 번 오큰수를 탐색한 것을 재탐색하지 않는 방법을 찾아야한다. 따라서 스택을 아이디어로 생각해낼 수 있다.

• 스택에 새로 들어오는 수가 top에 존재하는 수보다 크면 그 수는 오큰수가 된다.
• 오큰수를 구한 후 수열에서 오큰수가 존재하지 않는 숫자에 -1을 출력해야 한다.

문제 바로가기 💨 

위 예시를 차근차근 따라가 보자.

1. 스택이 비었으면 push()로 수열 배열의 첫 인덱스를 push()한다.
2. 다음 인덱스(1)가 들어올때, Top() 인덱스의 배열값이 더 작으면 pop()을 해준다.(오큰수 발견!)
3. 그 후 다음 인덱스(1)를 push한다.
4. 다음 인덱스(2)가 들어올때, Top() 인덱스의 배열값이 더 작지 않으면 다음 인덱스(3)를 append() 한다.
5. 다음 인덱스(3)이 들어올때, Top() 인덱스의 배열값(2)이 더 작으므로 pop()을 해준다.(마찬가지로 1도 pop())
6. 수열 길이만큼 끝나면 배열의 남은 자리에 -1을 넣어준다.

주의점 (시간 초과)

아래 코드는 빈 문자열 result를 생성하고, 반복문을 통해 ans 배열의 각 요소를 문자열로 변환하여 result에 더한다. 이 방식은 문자열을 추가할 때마다 새로운 문자열 객체를 생성하게 되므로, 큰 데이터셋 또는 많은 반복을 수행할 경우 메모리 사용과 실행 시간이 늘어날 수 있습니다. ( 파이썬의 문자열은 불변(immutable)이기 때문에, 각 += 연산마다 사실상 새로운 문자열이 생성되어 기존 문자열을 대체 )

result = ""
for i in range(n):
 result += str(ans[i])+" "
print(result)

따라서 아래와 같이 바꾼다. 이 코드는 sys.stdout.write를 사용하여 각 요소를 직접 표준 출력에 쓴다. sys.stdout.write 함수는 주어진 문자열을 출력 스트림에 바로 쓰므로, 중간에 문자열을 합치는 과정이 필요 없다.

즉, 이 방법은 중간에 추가적인 문자열 객체를 생성하지 않기 때문에, 출력할 내용이 많거나 성능이 중요한 경우에 자주 사용된다.

for i in range(n):
    sys.stdout.write(str(ans[i]) + " ")

슈도 코드 작성하기

N(수열 개수) A[](수열 배열) ans(정답 배열)
수열 배열 채우기
최초 스택 초기화
for(N만큼 반복) {
    while(스택 비어있지않고, 현재 수열 값이 top에 해당하는 수열보다 클 때까지){ //오큰수 발견
        pop()
        정답 배열에 오큰수를 현재 수열로 저장하기
    }
    현재 수열을 스택에 push()
}

while(스택이 빌 때까지){
    스택에 있는 인덱스의 정답 배열에 -1 저장하기
}
정답 배열 출력하기

 

JAVA

import java.util.*;
import java.lang.*;
import java.io.*;

class Main {
	public static void main (String[] args) throws java.lang.Exception {
	    BufferedReader bf = new BufferedReader(new InputStreamReader(System.in));
	    int n = Integer.parseInt(bf.readLine());
	    int[] A = new int[n]; // 수열 배열
	    int[] ans = new int[n]; // 정답 배열
	    
	    String[] str = bf.readLine().split(" ");
	    for(int i = 0; i < n; i++){
	        A[i] = Integer.parseInt(str[i]);
	    }
	    Stack<Integer> myStack = new Stack<>();
	    myStack.push(0); //처음에는 항상 스택이 비어 있으므로 최초 값을 push해 초기화
	    for (int i = 1; i < n; i++){
	        //스택 비어 있지 않음 && 현재 수열이 스택 top 인덱스가 가리키는 수열보다 큰 경우
	        while(!myStack.isEmpty() && A[myStack.peek()] < A[i]){
	            ans[myStack.pop()] = A[i]; // 정답 배열에 오큰수를 현재 수열로 저장
	        }
	        myStack.push(i); // 신규 데이터 push
	    }
	    while ( !myStack.isEmpty() ) {
	        // 반복문을 다 돌고 나왔는데 스택이 비어 있지 않다면 빌때까지
	        ans[myStack.pop()] = -1;
	    }
	    BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
	    for( int i = 0; i < n; i++){
	        bw.write(ans[i] + " ");
	    }
	    bw.write("\n");
	    bw.flush();
	}
}

PYTHON

import sys

n = int(input())
ans = [0] * n # 오큰수 정답 배열
A = list(map(int,input().split()))
myStack = []

for i in range(n):
    while myStack and A[myStack[-1]] < A[i]: # 오큰수 조건
        ans[myStack.pop()] = A[i]
    myStack.append(i)

while myStack:
    ans[myStack.pop()] = -1
    
for i in range(n):
    sys.stdout.write(str(ans[i])+" ")