Pythonで1回の二乗操作後に最大部分配列の合計を求めるプログラム

整数値を含む配列が与えられているとします。この配列に対して、「array[i] の値をその二乗（array[i] × array[i]）に置き換える」という操作を1回だけ行うことができます。操作後に得られる最大の部分配列（サブアレイ）の合計を返す必要があります。なお、部分配列は空であってはなりません。

例えば、入力が array = [4, 1, -2, -1] の場合、出力は 17 になります。

array[0] の値を二乗に置き換えると、配列は [16, 1, -2, -1] となります。このとき最大の合計を持つ部分配列は [16, 1] であり、その合計は 16 + 1 = 17 です。

解法のアプローチ

この問題は動的計画法（DP）を使うことで効率的に解けます。ポイントは次の2つのテーブルを並行して管理することです。

• dp1：まだ二乗操作を一度も使っていない状態での、各位置までの最大部分配列の合計（いわゆるカダネのアルゴリズム）
• dp2：すでに二乗操作を1回使用した状態での、各位置までの最大部分配列の合計

具体的な手順は以下の通りです。

• dp1 を負の無限大（-inf）を含む新しいリストとして初期化する
• dp2 を負の無限大（-inf）を含む新しいリストとして初期化する
• 配列内の各要素 num について以下を繰り返す
  • max（dp1の末尾要素 + num, num）を dp1 の末尾に追加する
  • max（dp1の末尾から2番目の要素 + num², num², dp2の末尾要素 + num）を dp2 の末尾に追加する
• dp2 の最大値を返す

ここで dp2 の遷移式には3つの選択肢があります。「直前まで操作済みの部分配列に現在の要素を足す」「現在の要素で新しく部分配列を始める（つまりここで初めて二乗操作を使う）」「操作前の状態（dp1）から現在の要素を二乗して遷移する」の3つです。これにより、二乗操作をどの位置で適用するのが最適かを自動的に判断できます。計算量は O(n)、追加のメモリもリスト2本分で済みます。

実装例

以下のPythonコードで実際の動作を確認してみましょう。

def solve(array):
    dp1 = [float('-inf')]
    dp2 = [float('-inf')]
    for num in array:
        dp1.append(max(dp1[-1] + num, num))
        dp2.append(max(dp1[-2] + num**2, num**2, dp2[-1] + num))
    return max(dp2)

print(solve([4, 1, -2, -1]))

入力

[4, 1, -2, -1]

出力

17