概要

社員を連続するグループに分割し、各グループの生産性（人数×能力値の合計）の合計を最大化する問題です。

考察

素朴なアプローチとして、全ての分割パターンを列挙すると、分割位置の選び方が \(O(2^N)\) 通りあり、\(N=5000\) では計算が不可能です。そこで動的計画法（DP）を用います。DP[r] を社員1から社員rまでの最適な生産性の合計と定義し、最後のグループの範囲 [l+1, r] を全探索することで状態遷移を行います。ただし、単純な実装では \(O(N^3)\) となり、\(N=5000\) では約1250億回の計算が必要でTLEとなります。そこで、累積和を用いてグループの合計計算を \(O(1)\) にし、全体を \(O(N^2)\) に高速化します。

アルゴリズム

1. 累積和配列 prefix_sum を事前計算します。prefix_sum[i] = \(P_1 + P_2 + \cdots + P_i\) となります。
2. DP配列を用意し、DP[0] = 0 で初期化します。
3. rを1からNまでループし、各rに対してlを0からr-1までループします：
   • グループ [l+1, r] のサイズは group_size = r - l
   • グループの能力値合計は group_sum = prefix_sum[r] - prefix_sum[l]
   • このグループを最後に追加した場合の生産性合計は candidate = DP[l] + group_size × group_sum
   • candidate が現在のDP[r]より大きければ更新します
4. DP[N]が求める答えとなります。

計算量

• 時間計算量: \(O(N^2)\)
• 空間計算量: \(O(N)\)

実装のポイント

• 累積和を事前計算することで、部分和の計算を \(O(1)\) にします。

• DP配列のインデックスは0-indexedで、DP[r]はr人目までの最適値を表します。

• 内側のループでlを0からr-1まで探索するため、二重ループで \(O(N^2)\) となりますが、\(N=5000\) では約1250万回の計算で実行可能です。

  ソースコード

import sys

def main():
    data = sys.stdin.read().split()
    n = int(data[0])
    P = list(map(int, data[1:1+n]))
    
    prefix_sum = [0] * (n + 1)
    for i in range(1, n + 1):
        prefix_sum[i] = prefix_sum[i - 1] + P[i - 1]
    
    dp = [0] * (n + 1)
    
    for r in range(1, n + 1):
        dp[r] = dp[r - 1] + P[r - 1]
        for l in range(r):
            group_size = r - l
            group_sum = prefix_sum[r] - prefix_sum[l]
            candidate = dp[l] + group_size * group_sum
            if candidate > dp[r]:
                dp[r] = candidate
                
    print(dp[n])

if __name__ == "__main__":
    main()

この解説は deepseekv3 によって生成されました。