概要

\(N\) 人の社員を番号が連続するグループに分割し、各グループの生産性（人数 × 能力値の合計）の総和を最大化する問題です。動的計画法（DP）を用いて解きます。

考察

問題の整理

社員 \(l, l+1, \ldots, r\) からなるグループの生産性は \((r - l + 1) \times \sum_{i=l}^{r} P_i\) です。社員を連続するグループに分割するので、分割点を決める問題と捉えられます。

具体例で理解する

例えば \(N = 4\), \(P = [3, -1, 2, 4]\) の場合を考えます。

• 全員1グループ: \(4 \times (3 + (-1) + 2 + 4) = 4 \times 8 = 32\)
• \(\{1,2\}, \{3,4\}\) に分割: \(2 \times (3+(-1)) + 2 \times (2+4) = 4 + 12 = 16\)
• \(\{1\}, \{2,3,4\}\) に分割: \(1 \times 3 + 3 \times 5 = 18\)

能力値の合計が正のグループはなるべく大きくまとめた方が得ですが、負の能力値を持つ社員が混ざるとグループを分けた方がよい場合もあり、最適な分割は自明ではありません。

DPの定式化

累積和を \(S[0] = 0\), \(S[i] = P_1 + P_2 + \cdots + P_i\) と定義します。

\(dp[i]\) を「社員 \(1\) から \(i\) までを最適に分割したときの生産性の合計の最大値」とします。初期値は \(dp[0] = 0\) です。

遷移は「最後のグループを社員 \(j+1\) から \(i\) とする」と考えて：

\[dp[i] = \max_{0 \leq j < i} \left( dp[j] + (i - j) \times (S[i] - S[j]) \right)\]

ここで \((i - j)\) はグループの人数、\((S[i] - S[j])\) はグループ内の能力値の合計です。

計算量の見積もり

\(dp[i]\) の計算で \(j\) を \(0\) から \(i-1\) まで走査するので、全体で \(O(N^2)\) の遷移があります。\(N \leq 5000\) なので、最大 \(\frac{5000 \times 5001}{2} \approx 1250\) 万回の演算となり、十分間に合います。

アルゴリズム

1. 入力を読み込み、累積和 \(S[0..N]\) を計算する。
2. \(dp[0] = 0\) と初期化する。
3. \(i = 1, 2, \ldots, N\) の順に、\(j = 0, 1, \ldots, i-1\) を全探索して次の遷移を計算する： $\(dp[i] = \max_{0 \leq j < i} \left( dp[j] + (i - j) \times (S[i] - S[j]) \right)\)$
4. \(dp[N]\) を出力する。

計算量

• 時間計算量: \(O(N^2)\)
• 空間計算量: \(O(N)\)（累積和配列と DP 配列）

実装のポイント

• 累積和の前計算: グループの能力値の合計を \(S[i] - S[j]\) で \(O(1)\) で求められるようにしておく。

• 内側ループの最適化: si = S[i] をループ外で変数に保持し、毎回配列アクセスすることを避ける。Python では定数倍が遅いため、こうした小さな工夫が重要です。

• オーバーフローへの注意: 問題文に「64 ビット符号付き整数型に収まる」とありますが、Python は多倍長整数を標準でサポートするため気にする必要はありません。

• Convex Hull Trick (CHT) の検討: 遷移式を展開すると \(dp[j] + j \cdot S[j] - i \cdot S[j] - j \cdot S[i] + i \cdot S[i]\) となり、\(-i \cdot S[j]\) と \(-j \cdot S[i]\) の両方に \(i, j\) の積が現れるため、単純な CHT の形には帰着しにくいです。\(N \leq 5000\) では \(O(N^2)\) で十分なので、素直な二重ループで解いています。

  ソースコード

import sys

def main():
    input_data = sys.stdin.buffer.read().split()
    N = int(input_data[0])
    P = [int(input_data[i + 1]) for i in range(N)]
    
    # prefix sum
    S = [0] * (N + 1)
    for i in range(N):
        S[i + 1] = S[i] + P[i]
    
    # dp[i] = maximum total productivity for employees 1..i
    # dp[0] = 0
    # dp[i] = max over j in [0, i-1] of dp[j] + (i - j) * (S[i] - S[j])
    #        = max over j of dp[j] + (i - j) * (S[i] - S[j])
    
    # Let's expand: dp[j] + (i - j) * (S[i] - S[j])
    # = dp[j] + i*S[i] - i*S[j] - j*S[i] + j*S[j]
    # = dp[j] + j*S[j] - i*S[j] - j*S[i] + i*S[i]
    # = (dp[j] + j*S[j]) - S[j]*i - j*S[i] + i*S[i]
    # = (dp[j] + j*S[j]) - i*S[j] - j*S[i] + i*S[i]
    
    # This is: for fixed i, maximize over j:
    # f(j) = dp[j] + j*S[j] - i*S[j] - j*S[i]
    # = (dp[j] + j*S[j]) - i*S[j] - j*S[i]
    
    # This doesn't factor into a simple CHT form easily because we have both -i*S[j] and -j*S[i].
    # Let's think of it as: for each j, the "line" in terms of variable x = i:
    # f_j(x) = (dp[j] + j*S[j]) - x*S[j] - j*S[i]
    # But S[i] depends on i=x, so this isn't a line in x alone.
    
    # Actually let me reconsider. We have:
    # val = dp[j] + (i-j)*(S[i]-S[j])
    # = dp[j] + i*S[i] - i*S[j] - j*S[i] + j*S[j]
    # 
    # Group by j-dependent terms vs i-dependent:
    # = [dp[j] + j*S[j]] + i*S[i] - i*S[j] - j*S[i]
    # = [dp[j] + j*S[j]] + i*(S[i] - S[j]) - j*S[i]
    # = [dp[j] + j*S[j] - j*S[i]] + i*(S[i] - S[j])
    
    # Hmm, with N=5000, O(N^2) should be fine: 25 million operations.
    
    dp = [0] * (N + 1)
    
    for i in range(1, N + 1):
        best = -float('inf')
        si = S[i]
        for j in range(i):
            # group from j+1 to i, size = i - j, sum = S[i] - S[j]
            val = dp[j] + (i - j) * (si - S[j])
            if val > best:
                best = val
        dp[i] = best
    
    print(dp[N])

main()

---

この解説は claude4.6opus-thinking によって生成されました。