Pythonでリストの隣接しない要素の最大合計を求めるプログラム

数値のリスト nums が与えられたとき、互いに隣接しない要素だけを選んだ場合の最大合計を返す関数を作ることを考えます。リストには 0 や負の数が含まれている場合もあります。

たとえば、入力が [3, 5, 7, 3, 6] のとき、出力は 16 になります。これは、3・7・6 を選ぶことで要素同士が隣接せず、合計 16 を達成できるためです。

解き方の手順

この問題は動的計画法(DP)の考え方を使うと、O(n) の計算量で効率よく解けます。手順は次のとおりです。

• リストの長さが 2 以下の場合は、max(nums) をそのまま返す
• noTake(現在の要素を選ばない場合の最大合計)を 0 で初期化する
• take(現在の要素を選ぶ場合の最大合計)を nums[0] で初期化する
• i を 1 からリストの末尾まで繰り返す
  • take := noTake + nums[i](現在の要素を選ぶなら、直前の要素は選べない)
  • noTake := max(noTake, take)(それまでの最良値を引き継ぐ)
• 最後に max(noTake, take) を返す

アルゴリズムのポイント

各要素について「選ぶ(take)」「選ばない(noTake)」という2つの状態を管理するのがポイントです。現在の要素を選ぶ場合は直前の要素が使えないため noTake + nums[i] となり、選ばない場合は直前までの最良値 max(noTake, take) を引き継ぎます。こうして配列全体を一度走査するだけで答えが求まります。

実装例

理解を深めるために、以下の実装を見てみましょう。

class Solution:
   def solve(self, nums):
      if len(nums) <= 2:
         return max(nums)
      noTake = 0
      take = nums[0]
      for i in range(1, len(nums)):
         take, noTake = noTake + nums[i], max(noTake, take)
      return max(noTake, take)

ob = Solution()
nums = [3, 5, 7, 3, 6]
print(ob.solve(nums))

入力

[3, 5, 7, 3, 6]

出力

16

このように、隣接しない要素の最大合計はシンプルなDPで線形時間に求められます。負の数が含まれる場合でも、何も選ばない状態(合計 0)が常に候補として残るため、正しく動作します。