Pythonで乗算演算を繰り返して最大スコアを求めるプログラムの実装方法

問題概要

サイズ n の配列 nums と、サイズ m の配列 multipliers(n ≥ m)が与えられます。配列は1始まりのインデックスで扱い、初期スコアは0です。ここで、ちょうど m 回の操作を行います。i 番目の操作(1始まり)では、次の処理を実行します。

• nums の先頭または末尾から値 x を1つ選ぶ
• multipliers[i] × x をスコアに加算する
• x を nums から取り除く

m 回の操作をすべて終えた時点での最大スコアを求めるのが目標です。

たとえば、nums = [5, 10, 15]、multipliers = [5, 3, 2] が入力された場合、出力は 115 になります。まず末尾の 15 を選んで 5 × 15 = 75 を獲得し、次に 10 を選んで 3 × 10 = 30 を加算(合計 105)、最後に残りの 5 を選んで 2 × 5 = 10 を加算し、合計 75 + 30 + 10 = 115 となるためです。

解法のアプローチ(動的計画法)

この問題は、貪欲法では最適解が保証できないため、動的計画法(DP)を用いて解きます。dp[i][j] は「nums の左端から i 個、右端から (m − 1 − j) 個をすでに取り除いた状態」における最大スコアを表します。このとき、次に使う乗数のインデックスは k = i + m − j − 1 となります。手順は以下の通りです。

• n := nums のサイズ、m := multipliers のサイズとする
• dp := サイズ m × (m+1) の2次元配列を作成し、すべて0で初期化する
• i を m − 1 から 0 まで逆順にループする
  • j を i から m − 1 までループする
    • k := i + m − j − 1
    • dp[i][j] = max(nums[i] × multipliers[k] + dp[i+1][j]、nums[j−m+n] × multipliers[k] + dp[i][j−1])
• dp[0] の最後の要素を返す

遷移の意味としては、「左端から値を取る場合」と「右端から値を取る場合」の2択を比較し、大きい方を採用します。i を逆順に処理することで、より操作回数が進んだ状態から順に確定させていける点がポイントです。

実装例

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

def solve(nums, multipliers):
   n, m = len(nums), len(multipliers)
   dp = [[0]*m for _ in range(m+1)]

   for i in reversed(range(m)):
      for j in range(i, m):
         k = i + m - j - 1
         dp[i][j] = max(nums[i] * multipliers[k] + dp[i+1][j], nums[j-m+n] * multipliers[k] + dp[i][j-1])

   return dp[0][-1]

nums = [5,10,15]
multipliers = [5,3,2]
print(solve(nums, multipliers))

入力

[5,10,15], [5,3,2]

出力

115

まとめ

本記事では、配列の両端から値を選んで乗算スコアを最大化する問題を、Pythonの動的計画法で解く方法を紹介しました。計算量は O(m²) となり、n が大きくても m に依存した効率的な求解が可能です。両端キューを使うタイプの最適化問題では、「左から何個取ったか」を状態として持つDPが有効であることを覚えておくと、応用問題にも役立ちます。