Pythonで最小グループの合計が最大になるようリストをk個に分割する方法

問題概要

数値のリスト nums と整数 k が与えられたとします。このリストを「連続する要素からなる」k個のグループに分割することを考えます。ここで「最小グループ」とは、各グループの合計値の中で最も小さいものを持つグループのことです。求めたいのは、その最小グループの合計値が取りうる最大値です。

例として、nums = [2, 6, 4, 5, 8]、k = 3 の場合を見てみましょう。リストを [2, 6]、[4, 5]、[8] の3つのグループに分割すると、それぞれの合計は 8、9、8 となり、最小グループの合計は 8 になります。どのように分割しても最小グループの合計が 8 を超えることはできないため、答えは 8 です。

解き方：二分探索（バイナリサーチ）

この問題は「答えを決め打ちして判定する」二分探索で効率的に解けます。ある値 target に対して「すべてのグループの合計を target 以上にできるか」を判定する関数 is_divisible() を用意し、答えの候補範囲を狭めていきます。target が達成可能なら、それより小さい値も必ず達成可能であるため、判定結果には単調性があり、二分探索が成立します。

アルゴリズムの手順

1. is_divisible(target) 関数を定義します。
2. target が 1 以下の場合は True を返します。
3. num_chunks := 0、current_sum := 0 で初期化します。
4. nums の各要素 x について以下を繰り返します。
   • current_sum に x を加算します。
   • current_sum が target 以上になったら、current_sum を 0 に戻し、num_chunks を 1 増やします。
   • num_chunks が k に達したら True を返します。
5. ループが終了したら False を返します。

メイン処理の手順

1. left := 1 とします。
2. right := (nums の全要素の合計) ÷ k + 1 とします。
3. left < right − 1 の間、以下を繰り返します。
   • mid := (left + right) ÷ 2 とします。
   • is_divisible(mid) が真なら left := mid、そうでなければ right := mid とします。
4. 最後に left を返します。

計算量は、判定関数が O(n)、二分探索が O(log S)（S は要素の総和）なので、全体で O(n log S) となり、全分割パターンを試す方法より大幅に高速です。

Pythonでの実装例

以下の実装を見ると、動作がより理解しやすくなります。

class Solution:
    def solve(self, nums, k):
        def is_divisible(target):
            if target <= 1:
                return True
            num_chunks = 0
            current_sum = 0
            for x in nums:
                current_sum += x
                if current_sum >= target:
                    current_sum = 0
                    num_chunks += 1
                    if num_chunks == k:
                        return True
            return False

        left = 1
        right = sum(nums) // k + 1
        while left < right - 1:
            mid = (left + right) // 2
            if is_divisible(mid):
                left = mid
            else:
                right = mid
        return left

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

入力

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

出力

8