けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AOJ 2877 水槽 (RUPC 2018 day1-D)

とてもよくある区間分割していくタイプの DP

問題概要

 N 個の正の整数  a_{0}, \dots, a_{N-1} がこの順に一列に並んでいる。これらの整数を  M 個以下の区間に分割したい。

このとき、「各区間の値の平均値の総和」として考えられる最小値を求めよ。

制約

  •  1 \le M \le N \le 500
  •  1 \le a_{i} \le 10^{9}

考えたこと

「各区間の値の平均値の総和」を区間分割方法のスコアを呼ぶことにする。「区間ごとに分割していく方法を最適化する」という問題は、次のような DP で解けることが多い (けんちょん本 5.6 節参照)。

  • dp[i] N 個の整数のうち最初の  i 個について、最適な区間分割をしたときのスコア

今回はこれに加えて「区間の個数を  M 個以下にする」という制約もあるので、次のようにする。

  • dp[i][k] N 個の整数のうち最初の  i 個について、 k 個の区間に分割することを考える。最適な区間分割をしたときのスコア

このとき、DP の遷移は次のように考えられる。これは、 i 個の要素を区間ごとに分割する仕切りのうち、最後の仕切りを入れる位置  j を場合分している。

chmin(dp[i][k], dp[j][k-1] + f(j, i))

ここで、 f(j, i) は、区間  \lbrack j, i) の値の平均値とする。 f(j, i) の値についてはあらかじめ前処理によって求めておこう。

以上の DP 遷移の計算量は  O(N^{2}M) となる。

コード

以下のコードでは、f[j][i] の計算に  O(N^{3}) をかけている (工夫次第で  O(N^{2}) へと改善可能) ので、全体の計算量は  O(N^{3} + N^{2}M) となっている。

#include <iostream>
#include <vector>
#include <iomanip>
using namespace std;
template<class T> inline bool chmax(T& a, T b) { if (a < b) { a = b; return 1; } return 0; }

const long long INF = 1LL<<60;
int main() {
    int N, M;
    cin >> N >> M;
    vector<long long> a(N);
    for (int i = 0; i < N; ++i) cin >> a[i];

    // 区間 [j, i) の平均値を前処理で求める
    vector<vector<double>> f(N+1, vector<double>(N+1));
    for (int i = 1; i <= N; ++i) {
        for (int j = 0; j < i; ++j) {
            double sum = 0;
            for (int k = j; k < i; ++k) sum += a[k];
            f[j][i] = sum / (i - j);
        }
    }

    // DP
    vector<vector<double>> dp(N+1, vector<double>(M+1, -INF));
    dp[0][0] = 0;
    for (int i = 0; i <= N; ++i) {
        for (int j = 0; j < i; ++j) {
            for (int k = 1; k <= M; ++k) {
                chmax(dp[i][k], dp[j][k-1] + f[j][i]);
            }
        }
    }
    double res = -INF;
    for (int m = 0; m <= M; ++m) chmax(res, dp[N][m]);
    cout << fixed << setprecision(10) << res << endl;
}