Pythonで要素を1つ移動してリストのパワーの最大値を求めるプログラム

問題の概要

N個の正の数からなるリスト nums が与えられていると仮定します。このリストから任意の1つの値を選び、それを(スワップではなく)移動させて好きな位置に挿入することができます。もちろん、どの要素も移動しないという選択も可能です。このとき、リストのパワーが取りうる最大値を求めるのがこの問題です。

ここでいうリストのパワーとは、すべてのインデックス i における (i + 1) × list[i] の総和として定義されます。

$$\displaystyle\sum\limits_{i=0}^{n-1} (i+1)\times list[i]$$

具体例

入力が nums = [6, 2, 3] の場合を考えてみましょう。値 6 を末尾へ移動すると、リストは [2, 3, 6] となり、パワーは次のように計算できます。

(2 × 1) + (3 × 2) + (6 × 3) = 2 + 6 + 18 = 26

したがって、この場合の出力は 26 となります。

解法のアプローチ

この問題は、累積和(prefix sum)を活用することで効率的に解けます。ある要素を位置 i から位置 j へ移動させたときのパワーの変化量を、累積和を使えば定数時間で求められるためです。

手順は以下の通りです。

• P := 値 0 で初期化したリスト(累積和を格納)
• base := 0(元のリストのパワー)
• A の各インデックス i と各値 x について、以下を実行します。
  • P の末尾に「P の最後の要素 + x」を追加する
  • base := base + (i + 1) × x
• ans := base(何も移動しない場合も候補とする)
• A の各インデックス i と各値 x について、以下を実行します。
  • j を 0 から A のサイズまで繰り返します。
    • ans := ans と (base + P[i] − P[j] − (i − j) × x) の大きい方
• ans を返す

式「base + P[i] − P[j] − (i − j) × x」は、「要素 x を位置 i から位置 j へ移動させた場合の新しいパワー」を表しています。P[i] − P[j] の部分は移動によって前後にずれる要素たちの寄与の変化を、−(i − j) × x の部分は移動する要素自体の重み((インデックス + 1))の変化を表しています。この式は、要素を前方へ移動する場合(j ≤ i)でも後方へ移動する場合(j > i)でも共通して成り立ちます。

実装例(Python)

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

class Solution:
    def solve(self, A):
        P = [0]
        base = 0
        for i, x in enumerate(A, 1):
            P.append(P[-1] + x)
            base += i * x
        ans = base
        for i, x in enumerate(A):
            for j in range(len(A) + 1):
                ans = max(ans, base + P[i] - P[j] - (i - j) * x)
        return ans
ob = Solution()
nums = [6, 2, 3]
print(ob.solve(nums))

入力

[6, 2, 3]

出力

26

計算量について

移動する要素と挿入先の位置のすべての組み合わせを試すため、時間計算量は O(n²) となります。もし毎回リストを実際に再構築してパワーを再計算する素朴な方法では O(n³) かかるところを、累積和のおかげで各移動の効果を O(1) で評価でき、大幅に高速化できます。必要な追加メモリは累積和リスト P のぶんだけで、空間計算量は O(n) です。