파이썬으로 반복 연산 후 가장 마지막에 0이 되는 인덱스 찾기

문제 정의

배열 A에 n개의 숫자가 있고, 또 다른 입력값 K가 주어진다고 가정해 보겠습니다. 우리는 주어진 연산을 반복 수행한 후 가장 마지막에 0으로 감소하는 인덱스를 찾아야 합니다.

연산의 규칙은 다음과 같습니다.

A[0]부터 A[N-1]까지 순서대로 각 요소를 A[i] = A[i] - K로 갱신합니다.
갱신 결과 A[i] < K가 되면 해당 값을 0으로 설정합니다.
한 번 0이 된 요소에는 더 이상 어떤 연산도 수행하지 않습니다.

이 연산을 모든 요소가 0이 될 때까지 반복하고, 가장 마지막에 0이 되는 인덱스를 반환하면 됩니다.

예시

입력이 A = [4, 3, 6, 8, 3, 10]이고 K = 4인 경우를 살펴보겠습니다. 이때 출력은 5입니다.

연산 1: A = {0, 0, 2, 4, 0, 6}
연산 2: A = {0, 0, 0, 0, 0, 2}
연산 3: A = {0, 0, 0, 0, 0, 0}

인덱스 5의 요소(값 10)가 세 번째 연산에서야 비로소 0이 되므로 정답은 5입니다.

접근 방법

실제로 연산을 하나씩 반복하며 시뮬레이션하는 것은 비효율적일 수 있습니다. 대신 수학적 아이디어를 활용하면 선형 시간 안에 해결할 수 있습니다.

핵심은 각 요소가 0이 되기 위해 필요한 연산 횟수를 미리 계산하는 것입니다. 요소 A[i]가 0이 되려면 ⌈A[i] / K⌉번의 연산이 필요하며, 이는 정수 연산으로 (A[i] + K - 1) // K와 동일합니다.

필요한 연산 횟수가 가장 큰 요소가 가장 늦게 0이 되고, 횟수가 같다면 더 뒤에 있는(인덱스가 큰) 요소가 매 연산에서 나중에 처리되므로 마지막에 0이 됩니다. 따라서 최대값을 찾되, 같은 값이면 뒤쪽 인덱스를 선택하면 됩니다.

알고리즘 단계는 다음과 같습니다.

n := 배열 A의 크기로 설정
idx := -1로 초기화
i를 0부터 n-1까지 반복하며 A[i] := (A[i] + k - 1) // k로 변환 (각 요소에 필요한 연산 횟수)
다시 i를 0부터 n-1까지 반복하며 A[i] >= x이면 x := A[i], idx := i로 갱신
idx 반환

구현 예제

다음 구현을 통해 더 잘 이해해 보겠습니다.

def search_index(A, k):
    n = len(A)
    idx = -1
    x = -10**9
    for i in range(n):
        A[i] = (A[i] + k - 1) // k
    for i in range(n):
        if (A[i] >= x):
            x = A[i]
            idx = i
    return idx

arr = [4, 3, 6, 8, 3, 10]
K = 4
print(search_index(arr, K))

입력

[4, 3, 6, 8, 3, 10], 4

출력

복잡도 분석

시간 복잡도: O(N) — 배열을 두 번 순회합니다.
공간 복잡도: O(1) — 추가 공간을 거의 사용하지 않습니다.