【CodeForces】908 D. New Year and Arbitrary Arrangement

【题目】Good Bye 2017 D. New Year and Arbitrary Arrangement

【题意】给定正整数k,pa,pb，初始有空字符串，每次有pa/(pa+pb)的可能在字符串末尾+a，有pb/(pa+pb)的可能在字符串末尾+b，求加到组成至少k对子序列“ab"时的期望子序列“ab”数。k<=1000，pa,pb<=10^6。

【算法】期望DP

【题解】主要问题在于字符串无限延伸，那么需要考虑记录前缀的关键量来为DP设置终止状态。

设f[i][j]表示前缀中有i个a和j个ab停止后的期望长度，A=pa/(pa+pb)，B=pb/(pa+pb)。

状态转移方程：f[i][j]=A*f[i+1][j]+B*f[i][i+j]

接下来解决两个问题：

1.终止状态：当i+j>=k时，再加一个b就会终止，期望为i+j+c，其中：

c=0*B+1*A*B+2*A^2*B+...+∞*A^∞*B

等差*等比数列，运用高中数学的错位相减法（特别的，A^∞=0），可以得到：

c=pa/pb

故有终止状态f[i][j]=i+j+pa/pb，i+j>=k。

2.初始状态：初始空字符串为f[0][0]，但是会发现f[0][0]会从f[0][0]本身转移，其原因是没有a时会无限加b，解决办法是初始状态设为f[1][0]。

#include<cstdio>
const int m=1e9+7,N=1010;
void gcd(int a,int b,int&x,int &y){
    if(!b){x=1;y=0;}
    else{gcd(b,a%b,y,x);y-=x*(a/b);}
}
int inv(int a){int x,y;gcd(a,m,x,y);return (x%m+m)%m;}
int f[N][N],k,pa,pb,A,B,C;
int main(){
    scanf("%d%d%d",&k,&pa,&pb);
    A=1ll*pa*inv(pa+pb)%m;B=(1-A+m)%m;C=1ll*pa*inv(pb)%m;
    for(int i=k;i>=1;i--)for(int j=k;j>=0;j--)
        f[i][j]=i+j>=k?(i+j+C)%m:(1ll*A*f[i+1][j]+1ll*B*f[i][i+j])%m;
    printf("%d",f[1][0]);
    return 0;
}