17298번: 오큰수 - Kotlin (스택)

17298번: 오큰수

문제

크기가 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)을 공백으로 구분해 출력한다.

 

나의 풀이

첫번째 풀이

 이 문제는 스택을 사용해서 수열의 원소를 하나씩 빼낸 다음, 남은 스택안에 있는 수열 중 현재 원소보다 크고 가장 왼쪽(즉 가장 먼저 발견되는 수)에 있는 수를 출력하는 문제이다. 나는 findlast와 엘비스 연산자를 사용해 문제를 해결하려고 하였지만 시간 초과가 발생하고 말았다.

import java.util.Stack

fun main() = with(System.`in`.bufferedReader()) {
    val n = readLine().toInt()
    val nge  = readLine().split(" ").map { it.toInt() }
    val stack = Stack<Int>()

    for (i in n - 1 downTo 0) {
        stack.add(nge[i])
    }

    while (stack.isNotEmpty()) {
        val now = stack.pop()
        print("${stack.findLast { now < it }?: -1} ")
    }
}

 

두번째 풀이

 스택을 너무 많이 탐색하는 것이 문제였다. 접근 방식이 문제였고 스택에 해당 원소를 직접 넣는게 아니라 index 번호를 넣어서 문제를 해결해야 했다. 스택에 peek 값과 i번째의 값을 비교를 하고 만약 i번째의 값이 더 크다면 해당 index에 i번째의 값을 넣게된다. 이 과정을 stack 안에 남아있는 원소가 없거나 peek 값이 i번째의 값보다 클때까지 반복한다. 이렇게 되면 stack안에 남아있는 index들은 자동적으로 오큰수를 찾지 못한 값들임으로 stack안의 index에 해당하는 nge들을 -1로 바꿔준다.

import java.util.Stack

fun main() = with(System.`in`.bufferedReader()) {
    val n = readLine().toInt()
    val nge  = readLine().split(" ").map { it.toInt() }.toIntArray()
    val stack = Stack<Int>()

    for(i in 0 until nge.size){
        while(stack.isNotEmpty() && nge[stack.peek()] < nge[i]){
            nge[stack.pop()] = nge[i]
        }
        stack.push(i)
    }

    while(stack.isNotEmpty()){
        nge[stack.pop()] = -1
    }

    val sb = StringBuilder()
    for(i in 0 until nge.size){
        sb.append("${nge[i]} ")
    }
    println(sb)
}