Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

Python으로 길이가 k인 겹치지 않는 세 개의 부분 리스트 최대 합 구하기

숫자로 이루어진 리스트 nums와 정수 k가 주어졌다고 가정해 보겠습니다. 이 문제의 목표는 리스트에서 길이가 k인 서로 겹치지 않는 세 개의 부분 리스트를 선택했을 때 얻을 수 있는 최대 합을 구하는 것입니다.

예를 들어 입력이 nums = [2, 2, 2, -6, 4, 4, 4, -8, 3, 3, 3], k = 3이라면 출력은 27이 됩니다. [2, 2, 2], [4, 4, 4], [3, 3, 3] 세 개의 부분 리스트를 선택하면 합이 6 + 12 + 9 = 27로 최대가 되기 때문입니다.

해결 접근 방법

이 문제는 단순히 모든 조합을 탐색하면 비효율적이지만, 누적 합(prefix sum)과 전방·후방 최댓값 배열을 활용하면 O(n) 시간에 효율적으로 해결할 수 있습니다. 단계별로 살펴보면 다음과 같습니다.

  • 누적 합 계산: P := [0]으로 초기화한 뒤, A의 각 원소 x에 대해 P의 끝에 P[-1] + x를 추가합니다.
  • 구간 합 계산: Q := [P[i + K] - P[i]] (i는 0부터 len(P) - K까지). 즉, Q에는 길이 K짜리 모든 윈도우의 합이 저장됩니다.
  • prefix 배열 생성: Q를 복사한 뒤 왼쪽부터 순회하며, 각 위치까지의 최댓값으로 갱신합니다. prefix[i]는 인덱스 i 이전 영역에서 얻을 수 있는 최대 구간합을 의미합니다.
  • suffix 배열 생성: 마찬가지로 Q를 복사한 뒤 오른쪽부터 순회하며, 각 위치 이후 영역의 최댓값으로 갱신합니다.
  • 세 구간 조합: 가운데 부분 리스트의 시작 인덱스 i(K ≤ i ≤ len(Q) - K - 1)에 대해 Q[i] + prefix[i - K] + suffix[i + K]를 계산합니다. 이는 왼쪽·가운데·오른쪽 세 구간의 합 중 가능한 최댓값입니다.
  • 계산된 값들 중 최댓값을 반환합니다.

Python 예제 코드

아래 구현을 통해 더 자세히 이해해 보겠습니다.

class Solution:
   def solve(self, A, K):
      P = [0]
      for x in A:
         P.append(P[-1] + x)
      Q = [P[i + K] - P[i] for i in range(len(P) - K)]
      prefix = Q[:]
      suffix = Q[:]
      for i in range(len(Q) - 1):
         prefix[i + 1] = max(prefix[i + 1], prefix[i])
         suffix[~(i + 1)] = max(suffix[~(i + 1)], suffix[~i])
      return max(Q[i] + prefix[i - K] + suffix[i + K] for i in range(K, len(Q) - K))
ob = Solution()
nums = [2, 2, 2, -6, 4, 4, 4, -8, 3, 3, 3]
k = 3
print(ob.solve(nums, k))

입력

[2, 2, 2, -6, 4, 4, 4, -8, 3, 3, 3], 3

출력

27

복잡도 분석

이 알고리즘은 누적 합 계산, prefix/suffix 배열 생성, 최댓값 탐색 과정이 모두 리스트를 한 번씩 순회하므로 시간 복잡도는 O(n)입니다. 추가로 사용되는 배열들의 크기가 입력 크기에 비례하므로 공간 복잡도 역시 O(n)입니다. 브루트포스 방식의 O(n³) 탐색에 비해 매우 효율적이며, 입력 크기가 커져도 안정적으로 동작합니다.