[백준 / 17298번] 오큰수 - (Java - Stack)

 

17298번: 오큰수

---

⭐️ 코드

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.util.ArrayDeque;
import java.util.StringTokenizer;

// 오큰수
public class Main {
    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        StringTokenizer st = new StringTokenizer(br.readLine());

        int N = Integer.parseInt(st.nextToken()); // A 수열 크기

        int[] A = new int[N]; // A 수열 [ 3, 5, 2, 7 ]

        // A 수열 초기값 입력
        st = new StringTokenizer(br.readLine()); // 3 5 2 7
        for (int i = 0; i < N; i++) {
            A[i] = Integer.parseInt(st.nextToken());
        }

        // stack에 A 배열의 index를 저장
        ArrayDeque<Integer> stack = new ArrayDeque<Integer>();
        for (int i = 0; i < A.length; i++) {
            while (!stack.isEmpty() && A[stack.peekLast()] < A[i]) {
                A[stack.pollLast()] = A[i];
            }
            // 스택이 비었거나 A[stack.peek()]가 A[i]보다 작은 경우 stack에 push
            stack.addLast(i);
        }

        // 끝나고 스택에 남은 index에 해당하는 A배열의 값은 자신보다 큰 수가 없기 떄문에 -1로 초기화
        while (!stack.isEmpty()) {
            A[stack.pollLast()] = -1;
        }

        // 출력
        for(int i =0; i<A.length; i++) {
            bw.write(A[i] + " ");
        }
        
        bw.flush();
    }
}

---

💡 문제 풀이

먼저 실패한 코드를 보자

// 실패 코드
import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.util.ArrayDeque;
import java.util.Arrays;
import java.util.StringTokenizer;

// 오큰수
public class Main {
    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        StringTokenizer st = new StringTokenizer(br.readLine());

        int N = Integer.parseInt(st.nextToken()); // A 수열 크기

        int[] A = new int[N]; // A 수열 [ 3, 5, 2, 7 ]

        // A 수열 초기값 입력
        st = new StringTokenizer(br.readLine()); // 3 5 2 7
        for(int i = 0; i < N; i++) {
            A[i] = Integer.parseInt(st.nextToken());
        }

        // 오큰수 찾기
        int[] res = new int[A.length];
        for(int i = 0; i < A.length; i++) {
            res[i] = findRightBiggerNum(i, A);
            bw.write(res[i] + " ");
        }
        bw.flush();

    }

    private static int findRightBiggerNum(int idx, int[] A) {
        // 가장 왼 쪽 찾기
        for(int i = idx+1; i<A.length; i++) {
            if (A[i] > A[idx]) {
                return A[i];
            }
        }
        return -1;
    }
}

첫 번째 코드는 findRightBiggerNum() 함수를 만들어서 실행하였다. 하지만 시간복잡도를 계산해 보면 O(N^2)의 시간이 걸리게 된다. 

항상 최악의 결과는 모두 -1이 출력되는 경우이기 때문이다.

 

문제는 읽어보면 충분히 이해를 할 수 있을 테니 완성된 코드에 대한 풀이를 해보자.

 

우리는 Stack을 이용해서 문제를 풀어야 한다. 하지만 스택에 값을 집어넣어서 비교하는 형태로 진행하면 안 된다. (그렇게 하면 시간초과 남)

 

// A 수열 초기값 입력
st = new StringTokenizer(br.readLine()); // 3 5 2 7
for (int i = 0; i < N; i++) {
    A[i] = Integer.parseInt(st.nextToken());
}

일단 먼저 수열의 초기값을 A라는 새로운 배열을 만들어 값을 집어넣고

 

// stack에 A 배열의 index를 저장
ArrayDeque<Integer> stack = new ArrayDeque<Integer>();
for (int i = 0; i < A.length; i++) {
    while (!stack.isEmpty() && A[stack.peekLast()] < A[i]) {
        A[stack.pollLast()] = A[i];
    }
    // 스택이 비었거나 A[stack.peek()]가 A[i]보다 작은 경우 stack에 push
    stack.addLast(i);
}

스택을 만들어 A[i]의 값과 A [stack.peekLast()]의 값을 비교한다.

 

포문 내부의 while문을 잘 살펴보면 스택이 비어있으면 스택에 index값인 i를 삽입해 준다. 

A [스택에 존재하는 index값]과 A [i] 값을 비교하게 되면 A [stack.peekLast()]의 값이 더 작다면 스택에서 그 값을 꺼내고 A [stack.peekLast()]그 자리에는 A [i]에 해당하는 값을 넣어준다. 

 

이해가 잘 안된다면 그림을 그려보면서 풀어보자

// 예제 input
4
9 5 4 8

1. 스택이 비어있으니 스택에 9를 집어넣는 것이 아닌 i의 값인 0을 집어넣는다.

    A배열 : [ 9, 5, 4, 8 ]

 

2. 스택이 비어있지 않고 ( 0 )

    A[1] = 5

    A[stack.peekLast()] = A [0] = 9 이므로

    9 < 5 이기 때문에 오큰수를 만족하지 않아서 stack에 index(1)를 집어넣고 다음 포문으로 넘어간다.

    A배열 : [ 9, 5, 4, 8 ]

 

3. 스택이 비어있지 않고 ( 0, 1 )

    A[2] = 4

    A[stack.peekLast()] = A [1] = 5 이므로

    5 < 4 이므로 오큰수를 만족하지 않아서 stack에 index(2)를 집어넣고 다음 포문으로 넘어간다.

    A배열 : [ 9, 5, 4, 8 ]

 

3. 스택이 비어있지 않고 ( 0, 1, 2 )

    A[3] = 8

    A[stack.peekLast()] = A [2] = 4 이므로

    4 < 8 이기 때문에 오큰수(while 조건)를 만족해서

    A[stack.pollLast()] = A [2]에 8을 집어넣어 준다.

    A배열 : [ 9, 5, 8, 8 ]

 

3-1. while문 계속 추가 실행

    스택이 비어있지 않고 ( 0, 1 )

    A[3] = 8

    A[stack.peekLast()] = A [1] = 5 이므로

    5 < 8 이기 때문에 오큰수(while 조건)를 만족해서

    A[stack.pollLast()] = A [1]에 8을 집어넣어 준다.

    A배열 : [ 9, 8, 8, 8 ]

 

3-2. while문 계속 추가 실행

    스택이 비어있지 않고 ( 0 )

    A[3] = 8

    A[stack.peekLast()] = A [0] = 9 이므로

    9 < 8 이기 때문에 오큰수(while 조건)를 만족하지 않아서 stack에 index(3)를 집어넣고 다음 포문으로 넘어간다.

    A배열 : [ 9, 8, 8, 8 ]

    stack[0, 3]

 

4. 스택에는 stack[0, 3]이 들어 있어서 0과 3번 인덱스에 있는 수는 오큰수를 만족하지 않는 9와 8이기 때문에 -1로 변경해 주면 끝이다

while (!stack.isEmpty()) {
    A[stack.pollLast()] = -1;
}

 

마지막으로 BufferedWriter를 쓰기 때문에 출력을 bw로 해주면 끝!

// 출력
for(int i =0; i<A.length; i++) {
    bw.write(A[i] + " ");
}
bw.flush();

---

👀 후기

Stack... 너 이렇게 쓰는 거였구나..