Pythonで配列を合計が等しいk個の連続する部分配列に分割できるか判定する方法

数値の配列 nums と整数 k が与えられたとき、nums を k個の連続する部分配列に分割して、それぞれの部分配列の要素の合計がすべて等しくなるようにできるかどうかを判定する問題を考えてみましょう。

問題の例

たとえば、入力が nums = [2, 5, 3, 4, 7]、k = 3 の場合を考えます。このとき、[(2, 5), (3, 4), (7)] のように3つに分割でき、各部分配列の合計はいずれも 7 で等しくなるため、出力は True になります。

解法のアプローチ

この問題は累積和（prefix sum）を使うことで効率的に解けます。全体の合計が k で割り切れない場合は即座に False を返せばよく、割り切れる場合は累積和を走査しながら、目標となる合計（全体の合計 ÷ k）に達する位置を順番に記録していきます。

具体的には、以下の手順で処理を行います。

1. n を nums のサイズとします。
2. cumul_sum として nums の全要素の累積和を計算します。
3. total_sum を cumul_sum[n - 1]（配列全体の合計）とします。
4. total_sum が k で割り切れない場合は False を返します。
5. count = 0、pos = -1 として初期化します。
6. i を 0 から n-1 までループします。
   • pos が -1 のときは sub = 0、それ以外のときは sub = cumul_sum[pos] とします。
   • cumul_sum[i] - sub が total_sum / k と等しければ、その位置で区切りが確定したことになるので、pos = i として count を1増やします。
   • そうでなく、cumul_sum[i] - cumul_sum[pos] が total_sum / k を超えた場合は、それ以降に有効な区切りが存在しないためループを抜けます。
7. 最後に、count が k と等しければ True、そうでなければ False を返します。

実装例

以下はPythonでの実装例です。

def solve(nums, k):
   n = len(nums)
   cumul_sum = [0 for i in range(n)]
   cumul_sum[0] = nums[0]
   for i in range(1, n):
       cumul_sum[i] = cumul_sum[i - 1] + nums[i]
   total_sum = cumul_sum[n - 1]
   if total_sum % k != 0:
       return False
   count = 0
   pos = -1
   for i in range(n):
       if pos == -1:
           sub = 0
       else:
           sub = cumul_sum[pos]
       if cumul_sum[i] - sub == total_sum / k:
           pos = i
           count += 1
       elif cumul_sum[i] - cumul_sum[pos] > total_sum / k:
           break
   return count == k

nums = [2, 5, 3, 4, 7]
k = 3
print(solve(nums, k))

入力

[2, 5, 3, 4, 7], 3

出力

True

計算量と補足

このアルゴリズムの時間計算量は O(n)、空間計算量も累積和の保存に O(n) 必要となります。配列を一度走査するだけで判定できるため、非常に効率的です。

なお、Pythonでは / 演算子は浮動小数点除算を行うため、厳密な整数比較が必要な場合は //（整数除算）を使用するとより安全です。また、この方法は「連続する部分配列」への分割のみを扱う点に注意してください。要素を並べ替えてよい場合は、貪欲法やバックトラッキングなど別のアプローチが必要になります。