Codeforces #293 Div.2 D: Ilya and Escalator

Codeforcesにpypyが導入されてるの知らなかった。典型dp問題。

問題はこちら。

問題

n人の人がエスカレーターの前に並んでいる。行列の先頭の人は時間1ごとにpの確率でエスカレーターに乗る。エスカレーターの長さは十分長いとして、t秒後にエスカレーターに乗っている人数の期待値を答えよ。

解法

動的計画法で解く。

dp[ある時間に][最大m人が乗っている] = (乗っている人数の期待値) としてループを回せばいい。

O(nt)でn,t ≦ 2000だったがpythonだとTLEとなってしまった。pypyだと以下のコードでも余裕を持ってACとなった。pypyが選択可能とは知らなかった。

n,p,t = raw_input().split()
n,p,t = int(n),float(p),int(t)
dp = [[0]*(n+1) for i in xrange(t+1)]
q = 1.-p
for i in xrange(1,t+1):
    for j in xrange(1,n+1):
        dp[i][j] += p * (dp[i-1][j-1]+1) + q * dp[i-1][j]
print "%.7f"%dp[t][n]

はじめpure pythonだと間に合わなかったので、C++で一旦解答を書き直した。

#include <iostream>
#include <cstdio>
#define REP(i,a,b) for(int i = (a); i < (b); i++)
using namespace std;

double dp[2001][2001];

int main(void){
    int n,t;
    double p;
    cin >> n >> p >> t;
    REP(i,1,t+1) {
        REP(j,1,n+1) {
            dp[i][j] = p * (dp[i-1][j-1]+1) + (1-p) * dp[i-1][j];
        }
    }
    printf("%.7f",dp[t][n]);
    return 0;
}