SRM 478 div1 hard: RandomApple

問題

TopCoder Statistics - Problem Statement

解法

N 個の箱とは別の箱 (another box) を「箱」と呼び, i 番目の箱を指すときは 「i 番目の box」とか言うことにします。

この問題では, j 種類目のりんごを箱から取り出す確率を直接求めるのではなく, 箱からもともと i 番目の box に入ってたりんごを取り出す確率 p[i] を求め, その後 i 番目の box の j 番目のりんごを取り出す確率を求めて掛け算する, という手法を取ります。

i 番目の box に入っている j 種類目のりんごの数を apple[i][j], i 番目の box に入っているりんごの合計数を total[i] とすると, i 番目の box の j 番目のりんごを取り出す確率は apple[i][j] / total[i] と簡単に求めることが出来ます。

問題は p[i] を求めることですが, まず dp[i][j] = (i 番目までの box を使った時, 全種類のりんごの合計数が j になるような場合の数) を求めます。
で, memo[j] = (i 番目の箱を必ず使うときに, 全種類のりんごの合計数が j になるような場合の数) というのを, dp[N][j] を用いてそれぞれ O(MAX) で求めます。MAX は箱に入っているりんごの合計数で, 最高でだいたい 200*50*50 です。

これは, 多分戻す DP ってやつだと思います…
最後に追加するかしないかを考えている box が i 番目の box であったとすると, dp の遷移は以下のようになっています。

合計りんご数が j+sum[i] になる, 最後に sum[i] を追加するような場合の数(つまり memo[j+sum[i]]) が正確にわかっていたとすると, dp の遷移から, memo[j] は, dp[N][j] - memo[j+sum[i]] となります。
これで memo[j] を求めることが出来, そこから p[i] は, (memo[j]/(2^N-1) * total[i]/j) の j = 0〜MAX の和で求められます。

const int MAX = 200*50*50;
int apple[55][55];
int total[55];
ll dp[2][MAX];
ll memo[MAX];
double p[55];

class RandomApple {
public:
    vector <double> theProbability(vector <string> hundred, vector <string> ten, vector <string> one) {
        int N = hundred.size(), K = hundred[0].size();
        memset(apple, 0, sizeof(apple));
        memset(total, 0, sizeof(total));
        for (int i = 0; i < N; i++) for (int j = 0; j < K; j++) {
            apple[i][j] = (hundred[i][j]-'0')*100 + (ten[i][j]-'0')*10 + (one[i][j]-'0');
        }
        for (int i = 0; i < N; i++) for (int j = 0; j < K; j++) {
            total[i] += apple[i][j];
        }
        memset(dp, 0, sizeof(dp));
        dp[0][0] = 1;
        for (int i = 0; i < N; i++) {
            int cur = i%2, tar = cur^1;
            for (int j = 0; j < MAX; j++) dp[tar][j] = 0;
            for (int j = 0; j < MAX; j++) {
                if (dp[cur][j] > 0) {
                    dp[tar][j] += dp[cur][j];
                    dp[tar][j+total[i]] += dp[cur][j];
                }
            }
        }
        for (int i = 0; i < N; i++) p[i] = 0;
        for (int i = 0; i < N; i++) {
            for (int j = 0; j < MAX; j++) memo[j] = dp[N%2][j];
            for (int j = MAX-1; j >= 0; j--) {
                if (memo[j] > 0) {
                    p[i] += 1.*memo[j]*total[i]/j;
                    memo[j-total[i]] -= memo[j];
                }
            }
            p[i] /= (1ll<<N)-1;
        }
        vector<double> ans(K, 0);
        for (int i = 0; i < N; i++) {
            for (int j = 0; j < K; j++) {
                ans[j] += p[i] * apple[i][j] / total[i];
            }
        }
        return ans;
    }
};