yukicoder : No.1706 Many Bus Stops (hard)

SRMでレート大幅減の後なのでしんどい。
https://yukicoder.me/problems/no/1706

問題

C個のバス停がある。
整数時刻にあるバス停にいるバスは、1/Cの確率で次の整数時刻までそのバス停にとどまり、その他の確率で等確率で他のバス停に1.5の時間をかけて移動する。

初期状態でM台のバスが1番目のバス停にいるとき、時刻N経った時に1番目のバス停にバスが1台以上いる確率を答えよ。

解法

1番目以外のバス停にいる確率はすべて等しいと考えると、整数時刻にバスがいる状態は以下のいずれかである。

• 1番目のバス停にいる
• 1番目以外のバス停にいる
• 1番目のバス停から1番目以外のバス停に移動中
• 1番目以外のバス停から1番目のバス停に移動中
• 1番目以外のバス停から別以外の1番目のバス停に移動中

あとは5*5の遷移行列を考え、N乗を取れば、1台のバスが1番のバス停に戻っている確率pが求まる。
最終的に求めたいのは、1台以上のバスがいる確率なので、1-(1-p)^Mを答えればよい。

const int MAT=5;
struct Mat { ll v[MAT][MAT]; Mat(){ZERO(v);};};
const ll mo=1000000007;
ll modpow(ll a, ll n = mo-2) {
	ll r=1;a%=mo;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

Mat mulmat(Mat& a,Mat& b,int n=MAT) {
	ll mo2=4*mo*mo;
	int x,y,z; Mat r;
	FOR(x,n) FOR(y,n) r.v[x][y]=0;
	FOR(x,n) FOR(z,n) FOR(y,n) {
		r.v[x][y] += a.v[x][z]*b.v[z][y];
		if(r.v[x][y]>mo2) r.v[x][y] -= mo2;
	}
	FOR(x,n) FOR(y,n) r.v[x][y]%=mo;
	return r;
}

Mat powmat(ll p,Mat a,int n=MAT) {
	int i,x,y; Mat r;
	FOR(x,n) FOR(y,n) r.v[x][y]=0;
	FOR(i,n) r.v[i][i]=1;
	while(p) {
		if(p%2) r=mulmat(r,a,n);
		a=mulmat(a,a,n);
		p>>=1;
	}
	return r;
}

int C,N,M;
void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>C>>N>>M;
	Mat A;
	// A,notA,A->notA,notA->A,notA->notA
	A.v[0][0]=modpow(C);
	A.v[2][0]=1+mo-modpow(C);
	A.v[1][1]=modpow(C);
	A.v[3][1]=modpow(C);
	A.v[4][1]=(1-2*modpow(C)+2*mo)%mo;
	
	A.v[1][2]=1;
	A.v[0][3]=1;
	A.v[1][4]=1;
	
	A=powmat(N,A);
	ll notA=1+mo-A.v[0][0];
	ll ret=1-modpow(notA,M)+mo;
	cout<<ret%mo<<endl;
	
}

まとめ

Easyの方が面倒だったかも？