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

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

yukicoder No.155 生放送とBGM

戻す DP シリーズのトレーニング

問題へのリンク

問題概要 (入力形式は多少改変)

正の整数  L と、 N 個の正の整数  S_{1}, \dots, S_{N} が与えられる。これらをランダムシャッフルする。

 \sum_{i = 1}^{K} S_{K} \ge L を満たす最小の  K を考える。この値の期待値を求めよ (要求精度は  10^{-9})。

制約

  •  1 \le N \le 50
  •  1 \le L \le 18000

考えたこと

各要素  S_{i} に対して、それが再生される確率を求めて合計していけばよい。

大まかな方針として、 S_{i} の前に  j 個の要素が来る場合について  S_{i} までの総和が  L 以下となる確率を求めることにする。この値を各  i, j に対して合計して  N で割ればよい。 S_{i} の前にくる  j 個の選び方は、 C(N-1, j) 通りある。そのうちの何通りについて、総和が  L - 1 以下となるかを考えよう。これは  N-1 個の要素についての部分和 DP をすることで求められる。

  • dp[ i ][ j ][ k ] :=  N-1 個の要素のうちの最初の i 個について、そのうちの j 個を選ぶ場合であって、総和が k となるものが何個あるか

さて、このままでは、各  i に対して、 N-1 個の要素のついての部分和問題を解くため、 O(N^{3}L) となってしまう。計算量的に厳しい。

戻す DP

 i に対して、 N 個の要素から  i 番目を除いた  N-1 個の要素についての解を得る方法として、「戻す DP」が有効。dp 遷移と、逆変換を考える。

dp 配列に値 v を追加するとき

配列 dp から配列 ndp への遷移は次のようになる

ndp[ j ][ k ] = dp[ j ][ k ] + dp[ j - 1 ][ k - v ]

その逆変換

dp[ j ][ k ] = ndp[ j ][ k ] - dp[ j - 1 ][ k - v ]

以上によって、計算量は  O(N^{2}L) に改善できる。

コード

#include <bits/stdc++.h>
using namespace std;

vector<vector<long long>> add(vector<vector<long long>> dp, int v) {
    auto res = dp;
    for (int j = 1; j < res.size(); ++j) {
        for (int k = v; k < res[0].size(); ++k) {
            res[j][k] += dp[j-1][k-v];
        }
    }
    return res;
}

vector<vector<long long>> sub(vector<vector<long long>> dp, int v) {
    auto res = dp;
    for (int j = 1; j < res.size(); ++j) {
        for (int k = v; k < res[0].size(); ++k) {
            res[j][k] -= res[j-1][k-v];
        }
    }
    return res;
}

double com(int n, int r) {
    double res = 1;
    for (int i = 0; i < r; ++i) {
        res *= (n-i);
        res /= (i+1);
    }
    return res;
}

int main() {
    int N, L;
    cin >> N >> L;
    L *= 60;
    vector<int> S(N);
    for (int i = 0; i < N; ++i) {
        string str;
        cin >> str;
        int m = (str[0]-'0')*10 + (str[1]-'0');
        int s = (str[3]-'0')*10 + (str[4]-'0');
        S[i] = m*60+s;
    }

    double res = 0;
    vector<vector<long long>> dp(N+1, vector<long long>(L+1, 0));
    dp[0][0] = 1;
    for (int i = 0; i < N; ++i) dp = add(dp, S[i]);
    for (int i = 0; i < N; ++i) {
        double tmp = 0;
        auto sdp = sub(dp, S[i]);
        for (int k = 0; k < N; ++k) {
            double num = 0;
            for (int j = 0; j <= L-1; ++j) num += sdp[k][j];
            tmp += num / com(N-1, k);
        }
        res += tmp / N;
    }
    cout << fixed << setprecision(10) << res << endl;
}