[프로그래머스/파이썬] Lv.0 가까운 1 찾기

문제 설명

arr : 정수 배열. 1 또는 0의 값

idx : 정수

idx보다 크면서 배열의 값이 1인 가장 작은 인덱스 찾아 반환하는 함수 완성하기.

 

제한사항

• 3 ≤ arr의 길이 ≤ 100'000
  • arr의 원소는 전부 1 또는 0입니다.

입출력 예

arr	idx	result
[0,0,0,1]	1	3
[1,0,0,1,0,0]	4	-1
[1,1,1,1,0]	3	3

 

입출력 예 설명

1. 1보다 크면서 원소가 1인 가장 작은 인덱스는 3입니다. 따라서 3을 return 합니다.
2. 4번 인덱스 이후에 1은 등장하지 않습니다. 따라서 -1을 return 합니다.
3. 3번 인덱스의 값이 1입니다. 따라서 3을 return 합니다.

나의 풀이

def solution(arr, idx):
    answer = 0
    for i in range(idx, len(arr)):
        if arr[i] == 1:
            answer = i
            break   # 중요 !!
        else:
            answer = -1
    return answer

 

처음 코드에서 break를 빼고 적어 3번 예제에서 계속 오류가 났다.

이유는 1인 i를 찾고 다시 0이 나오면 -1로 바뀌게 되는 것이다.

break를 넣어주면서 1을 찾으면 멈추게 해준다면 가장 작은 인덱스를 자연스레 찾게 된다.

다른 사람의 풀이

def solution(arr, idx):
    for i in range(idx, len(arr)):
        if arr[i] == 1:
            return i
    return -1

 

answer 을 쓰지 않아 더 간편한 코드가 되었다.

이는 i를 찾으면 바로 return하기 때문에 내 코드와 다르게 break가 필요 없다.