Pythonで左右の要素の合計が等しいインデックスを見つけるプログラム

リスト nums が与えられたとき、「あるインデックス i の左側にある要素の合計」と「右側にある要素の合計」が等しくなるような、最小のインデックス i を求めることを考えます。該当するインデックスが存在しない場合は -1 を返します。

例えば、入力が nums = [8,2,3,6,5,2,5,9,1,2] の場合、出力は 4 となります。これは、インデックス 4 の左側の要素 [8,2,3,6] の合計が 19、右側の要素 [2,5,9,1,2] の合計も 19 となり、両者が一致するためです。

解法のアプローチ

この問題を効率的に解くには、リスト全体の合計をあらかじめ計算しておき、左から順に走査しながら「右側の残りの合計」と「左側の累積合計」を比較していく方法が有効です。手順は以下の通りです。

• 変数 r に nums の全要素の合計を代入します
• 変数 l を 0 で初期化します
• enumerate() を使って、各インデックス i と値 x を順に処理します
  • r から x を引きます(これで現在位置より右側の合計になります)
  • r と l が等しければ、i を返します
  • l に x を加えます(次の位置のための左側の合計を更新)
• ループが終了しても条件を満たすインデックスが見つからなければ、-1 を返します

実装例

以下に Python での実装例を示します。

def solve(nums):
   r = sum(nums)
   l = 0
   for i,x in enumerate(nums):
      r -= x
      if r == l:
         return i
      l += x
   return -1

nums = [8,2,3,6,5,2,5,9,1,2]
print(solve(nums))

入力

[8,2,3,6,5,2,5,9,1,2]

出力

4

アルゴリズムのポイント

この手法では、各インデックスごとに左右の合計を毎回再計算する必要がないため、時間計算量は O(n)、追加のメモリ使用量は O(1) で済みます。全要素の合計を先に求めておき、走査中に右側の合計を減算しながら左側の合計を加算していくことで、たった 1 回のループで答えを導き出せるのが大きなポイントです。素朴な 2 重ループによる O(n²) の解法と比べて、大規模なデータでも高速に動作します。