Python에서 합이 가장 큰 k개의 부분 리스트를 찾고 합계를 오름차순으로 반환하는 프로그램

문제 소개

숫자로 이루어진 리스트 nums와 정수 k가 주어졌을 때, 합이 가장 큰 k개의 연속된 부분 리스트를 찾아 그 합들을 오름차순(비내림차순)으로 반환하는 프로그램을 작성해 보겠습니다.

예를 들어 입력이 nums = [2, 4, 5, -100, 12, -30, 6, -2, 6], k = 3이라면 결과는 [10, 11, 12]가 됩니다. 합이 가장 큰 세 개의 부분 리스트는 각각 다음과 같습니다.

[6, -2, 6] → 합 10
[2, 4, 5] → 합 11
[12] → 합 12

해결 전략: 누적 합과 최소 힙

이 문제는 누적 합(prefix sum)과 최소 힙(min-heap)을 조합하면 깔끔하게 해결할 수 있습니다. 파이썬의 heapq 모듈은 최소 힙만 지원하기 때문에, 구간 합을 음수로 바꿔 저장하면 가장 큰 합부터 차례대로 꺼낼 수 있다는 점을 활용합니다.

알고리즘의 단계별 흐름은 다음과 같습니다.

nums의 길이보다 1만큼 큰 누적 합 배열 ps를 만들고 모든 값을 0으로 초기화합니다.
리스트를 순회하며 ps[i + 1] = nums[i] + ps[i] 형태로 누적 합을 채웁니다.
빈 힙 hp를 생성합니다.
모든 인덱스 쌍 (i, j)에 대해 구간 합인 ps[j] - ps[i]를 음수로 변환해 힙에 삽입합니다.
힙에서 k개의 원소를 꺼내 부호를 되돌린 뒤 역순으로 배치하여 오름차순 결과를 만듭니다.

구현 예제

from heapq import heappop, heappush

class Solution:
    def solve(self, nums, k):
        # 1단계: 누적 합 배열 초기화
        ps = [0 for _ in range(len(nums) + 1)]

        # 2단계: 누적 합 계산
        for i, v in enumerate(nums):
            ps[i + 1] = v + ps[i]

        # 3~4단계: 모든 구간 합을 음수로 힙에 삽입
        hp = []
        for i in range(len(ps)):
            for j in range(i + 1, len(ps)):
                heappush(hp, -(ps[j] - ps[i]))

        # 5단계: k개를 꺼내 오름차순으로 반환
        return list(reversed([-heappop(hp) for _ in range(k)]))

ob = Solution()
nums = [2, 4, 5, -100, 12, -30, 6, -2, 6]
k = 3
print(ob.solve(nums, k))

입력

nums = [2, 4, 5, -100, 12, -30, 6, -2, 6]
k = 3

출력

[10, 11, 12]

복잡도 분석

리스트의 길이를 n이라고 하면, 가능한 모든 구간의 수는 약 n²/2개입니다. 각 구간 합을 힙에 삽입하는 데 로그 시간이 걸리므로 전체 시간 복잡도는 O(n² log n)이며, 모든 구간 합을 힙에 저장하므로 공간 복잡도는 O(n²)입니다.

k가 전체 구간 수보다 훨씬 작은 경우에는 힙의 크기를 k로 제한하는 방식으로 메모리 사용량을 줄일 수 있으므로, 입력 크기에 따라 적절한 최적화를 고려해 보는 것도 좋습니다.