Pythonで要素の削除から得られる最大ポイントを求めるプログラム

問題の概要

正の整数のリストが与えられているとします。ここで、「同じ値が連続する長さ t の部分リスト」を一度に取り除くことができ、その際に t × t のポイントを獲得できるものとします。この操作は、リストが空になるまで何度でも繰り返し行えます。求めたいのは、最適な順序で削除を行ったときに獲得できるポイントの最大値です。

たとえば、入力が nums = [4, 4, 6, 4, 4] の場合、出力は 17 になります。

この結果が得られる流れを見てみましょう。まず、長さ1の「6」を取り除いて 1 × 1 = 1 ポイントを獲得します。すると残りのリストは [4, 4, 4, 4] となり、この長さ4の連続した「4」をまとめて取り除けば 4 × 4 = 16 ポイントを獲得できます。合計で 17 ポイントというわけです。

このように、間に挟まれた要素を先に除去することで、同じ値の要素を結合させ、より大きなスコアにつなげるのがポイントになります。

解法のアプローチ：区間DP（メモ化再帰）

この問題は、区間を単位とする動的計画法で効率よく解けます。dp(left, right, t) を「区間 [left, right] の左側に、nums[left] と同じ値の要素がすでに t 個つながっている状態における最大ポイント」と定義します。手順は以下の通りです。

• 関数 dp(left, right, t) を定義します。
• left > right の場合（区間が空の場合）は 0 を返します。
• num := nums[left] とし、left2 := left で初期化します。
• left2 < right かつ nums[left2 + 1] == num の間 left2 を進め、先頭から続く同じ値の個数を数えます。
• t := t + (left2 - left + 1)、left := left2 + 1 と更新します。
• points := t ** 2 + dp(left, right, 0) として、現在の連続ブロックをまとめて消去したケースのスコアを計算します。
• mid を left から right まで走査し、nums[mid] == num となる位置があれば、points = max(points, dp(left, mid - 1, 0) + dp(mid, right, t)) で更新します。これは「間の要素を先に処理してから、後ろの同じ値と結合する」戦略に相当します。
• 最終的な points を返します。

メイン側では dp(0, len(nums) - 1, 0) を呼び出して結果を出力します。

実装例

class Solution:
   def solve(self, nums):
      def dp(left, right, t):
         if left > right:
            return 0
         num = nums[left]
         left2 = left
         while left2 < right and nums[left2 + 1] == num:
            left2 += 1
         t += left2 - left + 1
         left = left2 + 1
         points = t ** 2 + dp(left, right, 0)
         for mid in range(left, right + 1):
            if nums[mid] == num:
               points = max(points, dp(left, mid - 1, 0) + dp(mid, right, t))
         return points
      return dp(0, len(nums) - 1, 0)

ob1 = Solution()
print(ob1.solve([4, 4, 6, 4, 4]))

入力

[4, 4, 6, 4, 4]

出力

17

まとめ

このアルゴリズムの本質は、「同じ値をすぐに消すか、それとも間の障害物を先に取り除いて後で大きな塊として消すか」という選択を、再帰的にすべて試す点にあります。実運用では functools.lru_cache などでメモ化を組み合わせると、重複する部分問題の計算を避けられ、大幅な高速化が期待できます。