Pythonで合計が最大のk個の部分リストを見つけ、合計を昇順で返す方法

問題概要

数値のリスト nums と整数 k が与えられたとき、合計が最も大きくなる k 個の部分リスト（連続する要素からなる区間）を見つけ、その合計を昇順（非減少順）で返すことを考えます。

たとえば、nums = [2, 4, 5, -100, 12, -30, 6, -2, 6]、k = 3 の場合、出力は [10, 11, 12] になります。これは、合計が大きい上位3つの部分リストがそれぞれ [6, -2, 6]（合計10）、[2, 4, 5]（合計11）、[12]（合計12）であるためです。

解法のアプローチ

この問題は「累積和（prefix sum）」と「ヒープ」を組み合わせることで解くことができます。手順は以下の通りです。

• nums のサイズ + 1 の長さを持つ累積和配列 ps を作成し、すべて 0 で初期化します。
• nums の各要素 v に対して、ps[i + 1] = v + ps[i] とし、累積和を格納していきます。
• 空のヒープ hp を用意します。
• すべてのインデックスのペア (i, j)（i < j）について、部分リスト nums[i..j-1] の合計である ps[j] - ps[i] を負の値にしてヒープへ挿入します。
• ヒープから要素を k 回取り出し、符号を元に戻した上で昇順に並べ替えて返します。

Python の heapq モジュールは最小ヒープのみを提供しているため、合計をあえて負の値として格納することで、実質的に最大ヒープとして動作させています。これにより、ヒープから取り出した順に「合計の大きい順」の値が得られます。

実装例

from heapq import heappop, heappush

class Solution:
    def solve(self, nums, k):
        # 累積和の計算
        ps = [0] * (len(nums) + 1)
        for i, v in enumerate(nums):
            ps[i + 1] = v + ps[i]

        hp = []
        # すべての部分リストの合計を負の値としてヒープに格納
        for i in range(len(ps)):
            for j in range(i + 1, len(ps)):
                heappush(hp, -(ps[j] - ps[i]))

        # 上位k個を取り出し、昇順に並べ替えて返す
        return sorted(-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))

入力

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

出力

[10, 11, 12]

計算量について

この実装では、考えられるすべての部分リスト（O(n²) 個）をヒープに挿入するため、時間計算量は O(n² log n)、空間計算量は O(n²) となります。リストのサイズ n が大きい場合は、ヒープに保持する要素数を k 個までに制限することで、メモリ使用量を O(k) に抑える最適化も可能です。