코딩테스트

가까운 1 찾기

kchabin 2025. 2. 10. 10:40

원소가 0, 1뿐인 정수 배열 arr과 인덱스 값 idx가 주어졌을 때, idx보다 크면서 배열의 값이 1인 가장 작은 인덱스를 반환하는 solution 함수를 완성해야 한다. idx보다 큰 인덱스의 값이 1인 경우가 없으면 -1을 반환해야 한다.

 

입출력 예

[0, 0, 0, 1] 1 3
[1, 0, 0, 1, 0, 0] 4 -1
[1, 1, 1, 1, 0] 3 3

 

두 번째 예시 같은 경우, idx=4가 주어지는데 그 이후 인덱스 5의 값이 0이고, 인덱스 4의 값도 0이라 -1을 반환해야 한다.

세번째 예시는 인덱스 3의 값이 1이지만 그 이후 인덱스 4는 0이고 이보다 더 큰 인덱스가 더 이상 존재하지 않는다. 이 경우 3을 리턴한다. 

근데 문제에서는 idx보다 '크면서' 조건을 만족하는 인덱스를 리턴해야 하는데, idx 자체를 리턴해도 되는걸까?

 

코드는 작성하긴 굉장히 간단한데, break을 추가해주지 않으니까 이후 인덱스가 0이면 answer값이 -1로 바뀌게 되는 문제가 있었다. 

class Solution {
    fun solution(arr: IntArray, idx: Int): Int {
        var answer: Int =0
        var len = arr.size
        for(i in idx until len) {
            if(i>=idx && arr[i]==1){
                answer = i
                break
            } 
            else answer = -1
            
             
        }
        return answer
    }
}

 

다른 사람 풀이 중에 진짜 대박 머리 좋다고 생각한 풀이.. firstOfNull을 사용해서 가장 처음 나오는 인덱스, 즉 idx 이후 가장 작은 인덱스를 뽑으면서 동시에 조건을 만족하는 값이 없으면 null을 출력해서 엘비스 연산자(?:)로 간단하게 답을 도출할 수 있는 코드가 있었다.

코틀린이 개인적으로  파이썬보다 코딩테스트에서 유용한 메서드가 많다고 생각하는데 문제는 너무 많아서 맨날 외우고 까먹는다는것...

fun solution(arr: IntArray, idx: Int): Int = (idx until arr.size).firstOrNull { arr[it] == 1 } ?: -1
}