Python으로 k번 반복된 리스트의 최대 연속 부분합 구하기

숫자 리스트 nums와 정수 k가 주어졌다고 가정해 봅시다. 여기서 k는 nums를 k번 이어 붙여 만든 긴 리스트를 의미하며, 우리는 그 안에서 합이 가장 큰 연속된 부분 리스트(연속 부분 배열)의 합을 구해야 합니다.

예를 들어 입력이 nums = [2, 4, 5, -4], k = 1이라면 출력은 11이 됩니다. [2, 4, 5]처럼 앞의 세 원소를 선택했을 때 합이 11로 가장 크기 때문입니다.

접근 방법

이 문제는 유명한 카데인 알고리즘(Kadane's Algorithm)을 확장한 형태로 해결할 수 있습니다. 핵심 아이디어는 다음과 같습니다.

k가 1 또는 2일 경우, 리스트를 최대 두 번 순회하며 누적합을 추적하면 최대 부분합을 찾을 수 있습니다.
k가 3 이상일 경우, 가운데에 반복되는 완전한 리스트 복사본들은 각각 전체 합 sum(nums)만큼씩 기여합니다. 전체 합이 양수라면 (k − 2)를 곱해 더하고, 음수라면 더하지 않는 것이 이득입니다.

알고리즘 단계

s, ans, lo를 모두 0으로 초기화합니다.
0부터 min(k, 2)까지의 범위 동안 반복합니다.
- nums의 각 원소 x에 대해 다음을 수행합니다.
  - s := s + x (누적합 갱신)
  - lo := min(lo, s) (지금까지의 최소 누적합 기록)
  - ans := max(ans, s − lo) (최대 부분합 갱신)
최종적으로 ans + max(0, sum(nums)) × max(0, k − 2)를 반환합니다.

동작 원리

변수 s는 현재까지의 누적합을, lo는 지금까지 등장한 누적합 중 최솟값을 저장합니다. 임의의 시점에서 최대 부분합은 '현재 누적합 − 지금까지의 최소 누적합'으로 표현할 수 있으므로, 매 단계마다 s − lo의 최댓값을 ans에 기록하면 됩니다. 이는 카데인 알고리즘의 전형적인 구현 방식입니다.

k가 2보다 클 때는 리스트 전체가 여러 번 반복되므로, 시작 부분과 끝 부분을 제외한 중간의 완전한 복사본들은 각각 sum(nums)만큼 합을 더해 줍니다. 전체 합이 음수라면 오히려 손해이므로 max(0, ...) 처리를 통해 기여분을 0으로 만들어 줍니다. 덕분에 실제로 거대한 리스트를 메모리에 만들지 않고도 O(n) 시간에 답을 구할 수 있습니다.

구현 예제

아래 파이썬 코드로 위 알고리즘을 구현할 수 있습니다.

class Solution:
    def solve(self, nums, k):
        s = ans = lo = 0
        for _ in range(min(k, 2)):
            for x in nums:
                s += x
                lo = min(lo, s)
            ans = max(ans, s - lo)
        return ans + max(0, sum(nums)) * max(0, (k - 2))

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

입력

[2, 4, 5, -4], 1