파이썬으로 K 접두사 문제 풀기: 누적 합이 K 이하가 되는 최대 인덱스 찾기

문제 정의

숫자 리스트 nums와 정수 k가 주어졌을 때, nums[0] + nums[1] + ... + nums[i] ≤ k를 만족하는 최대 인덱스 i를 찾아야 합니다. 만약 조건을 만족하는 인덱스가 존재하지 않는다면 -1을 반환합니다.

예시로 이해하기

예를 들어 nums = [4, -7, 5, 2, 6], k = 5라고 가정해 보겠습니다. 이 경우 출력값은 3입니다.

그 이유는 다음과 같습니다. 인덱스 3까지의 누적 합은 4 + (-7) + 5 + 2 = 4로 k보다 작거나 같습니다. 하지만 마지막 요소 6까지 더하면 합이 10이 되어 k를 초과하게 됩니다. 따라서 정답은 인덱스 3입니다.

해결 접근 방법

이 문제는 다음 두 단계를 통해 효율적으로 해결할 수 있습니다.

1. 누적 합 계산: 인덱스 1부터 리스트 끝까지 순회하면서 각 위치에 이전 값까지의 누적 합을 저장합니다. 즉, nums[i] = nums[i] + nums[i-1] 연산을 수행합니다.
2. 역방향 탐색: 리스트의 끝에서부터 시작 부분까지 거꾸로 순회하면서 nums[i] ≤ k를 처음으로 만족하는 인덱스를 찾아 반환합니다. 뒤에서부터 탐색하기 때문에 조건을 만족하는 가장 큰 인덱스를 바로 얻을 수 있습니다.

모든 인덱스를 확인했음에도 조건을 만족하는 값이 없다면 -1을 반환합니다.

구현 예제

class Solution:
    def solve(self, nums, k):
        for i in range(1, len(nums)):
            nums[i] += nums[i-1]
        for i in range(len(nums)-1, -1, -1):
            if nums[i] <= k:
                return i
        return -1

ob = Solution()
nums = [4, -7, 5, 2, 6]
k = 5
print(ob.solve(nums, k))

입력

[4, -7, 5, 2, 6], 5

출력

3

복잡도 분석

이 알고리즘의 시간 복잡도는 리스트를 두 번 순회하므로 O(n)입니다. 또한 별도의 추가 배열을 사용하지 않고 기존 리스트를 그대로 활용하기 때문에 공간 복잡도는 O(1)입니다.