概要

\(N\) 人の社員を連続するいくつかのグループに分割し、各グループの「人数 \(\times\) 能力値の合計」の総和を最大化する問題です。累積和と動的計画法（DP）を組み合わせることで、\(O(N^2)\) で効率的に解くことができます。

考察

連続する要素をいくつかの区間に分割して最適なスコアを求める問題では、「先頭から \(i\) 番目までを分割したときの最大値」を状態とする1次元の動的計画法（DP）が有効です。

まず、グループ内の「能力値の合計」を素早く計算するために累積和を使います。 社員の能力値の配列を \(P\) とし、その累積和を \(S\) とします。すなわち \(S[i] = \sum_{k=1}^i P_k\) です。 このとき、社員 \(j+1\) から \(i\) までのグループの能力値の合計は \(S[i] - S[j]\) と \(O(1)\) で求めることができます。

次に、DPの遷移を考えます。 \(dp[i]\) を「先頭から \(i\) 人目までを最適なグループに分割したときの生産性の最大値」と定義します。 最後のグループが社員 \(j+1\) から \(i\) まで（人数は \(i - j\) 人）であると仮定すると、このグループの生産性は以下のようになります。 $\( \text{生産性} = (i - j) \times (S[i] - S[j]) \)$

したがって、\(dp[i]\) を求めるには、直前の区切りの位置 \(j\) （\(0 \leq j < i\)）をすべて試し、最もスコアが高くなるものを選べばよいことになります。 $\( dp[i] = \max_{0 \leq j < i} \{ dp[j] + (i - j) \times (S[i] - S[j]) \} \)$

状態数が \(N\)、各状態での遷移（\(j\) の探索）が \(O(N)\) なので、全体の計算量は \(O(N^2)\) となります。\(N \leq 5000\) という制約のもとでは、\(O(N^2)\) のアルゴリズムで十分に制限時間に間に合います。

アルゴリズム

1. 累積和の構築 能力値 \(P\) の累積和配列 \(S\) を作成します。\(S[0] = 0\) とし、\(S[i] = S[i-1] + P_{i-1}\) と計算します。
2. DP配列の初期化 \(dp\) 配列を長さ \(N+1\) で用意し、\(0\) で初期化します。
3. DPの遷移 \(i = 1, 2, \ldots, N\) の順に以下の計算を行います。 $\( dp[i] = \max_{0 \leq j < i} \{ dp[j] + (i - j) \times (S[i] - S[j]) \} \)$
4. 答えの出力 最終的な答えは \(dp[N]\) に格納されます。

計算量

• 時間計算量: \(O(N^2)\) \(N\) 個の状態それぞれについて、最大 \(N\) 回の遷移を計算するためです。
• 空間計算量: \(O(N)\) 累積和配列 \(S\) と DP配列 \(dp\) などを保持するため、要素数 \(N\) に比例したメモリを使用します。

実装のポイント

Pythonなどのスクリプト言語で \(O(N^2)\) の二重ループを素朴に実装すると、定数倍が重くなり実行時間制限（TLE）に引っかかる可能性があります。そこで、計算式を展開してループ内の処理を軽くする工夫が有効です。

遷移式を展開すると以下のようになります。 $\( (i - j) \times (S[i] - S[j]) = i \cdot S[i] - i \cdot S[j] - j \cdot S[i] + j \cdot S[j] \)$

ここで、\(i \cdot S[i]\) は \(j\) に依存しないため \(\max\) の外に出すことができます。 $\( dp[i] = \max_{0 \leq j < i} \{ (dp[j] + j \cdot S[j]) - i \cdot S[j] - S[i] \cdot j \} + i \cdot S[i] \)$

実装では、あらかじめ \(A[j] = dp[j] + j \cdot S[j]\) を配列として保持しておくことで、ループ内の計算を最小限に抑えています。さらに、Pythonのリスト内包表記と組み込みの max() 関数を使用することで、C言語レベルで高速にループ処理が実行され、安全に制限時間内に通すことができます。

ソースコード

import sys

def solve():
    input = sys.stdin.read
    data = input().split()
    if not data:
        return
    N = int(data[0])
    P = [int(x) for x in data[1:]]
    
    S = [0] * (N + 1)
    for i in range(N):
        S[i+1] = S[i] + P[i]
        
    dp = [0] * (N + 1)
    A = [0] * (N + 1)
    B = S
    
    for i in range(1, N + 1):
        Si = B[i]
        dp_i = max([A[j] - i * B[j] - Si * j for j in range(i)]) + i * Si
        dp[i] = dp_i
        A[i] = dp_i + i * Si
        
    print(dp[N])

if __name__ == '__main__':
    solve()

---

この解説は gemini-3.1-pro-thinking によって生成されました。