Submission #1287828

#	Time	Username	Problem	Language	Result	Execution time	Memory
1287828		KhoaDuy	Star Trek (CEOI20_startrek)	C++20	100 / 100	94 ms	30864 KiB

startrek

#include <bits/stdc++.h>
using namespace std;
#define endl '\n'
const int mod=1e9+7;
int binpow(int n,long long k){
	int curr=1;
	for(;k;k>>=1){
		if(k&1){
			curr=(1LL*curr*n)%mod;
		}
		n=(1LL*n*n%mod);
	}
	return curr;
}
int x=0,y=0,z=0;
const int MAXN=1e5;
vector<vector<int>> graph(MAXN+1);
bool res[MAXN+1];
int sz[MAXN+1];
struct bruh{
	int x=0,y=0,z=0;
};
bruh DP[MAXN+1];
void setup(int u,int p){
	res[u]=false;
	sz[u]=1;
	for(int v:graph[u]){
		if(v!=p){
			setup(v,u);
			sz[u]+=sz[v];
			if(!res[v]){
				res[u]=true;
			}
		}
	}
	int child=0,lose=0;
	for(int v:graph[u]){
		if(v==p){
			continue;
		}
		child++;
		if(res[v]){
			lose++;
		}
	}
	if(lose==child){
		DP[u].y++;
	}
	else{
		DP[u].x++;
	}
	for(int v:graph[u]){
		if(v==p){
			continue;
		}
		if(res[v]){
			lose--;
		}
		if(lose==child-1){
			DP[u].x+=(sz[v]-DP[v].x-DP[v].y-DP[v].z);
			DP[u].y+=DP[v].z;
			DP[u].z+=DP[v].y;
		}
		else{
			DP[u].x+=sz[v];
		}
		if(res[v]){
			lose++;
		}
	}
}
#define x1 brehx1
#define y1 brehy1
#define z1 brehz1
int x1,y1,z1;
int wcnt=0;
void DFS(int u,int p){
	res[u]=false;
	sz[u]=1;
	for(int v:graph[u]){
		sz[u]+=sz[v];
		if(!res[v]){
			res[u]=true;
		}
	}
	if(res[u]){
		wcnt++;
	}
	vector<int> candi;
	DP[u].x=0,DP[u].y=0,DP[u].z=0;
	int child=0,lose=0;
	for(int v:graph[u]){
		child++;
		if(res[v]){
			lose++;
		}
		else{
			candi.push_back(v);
		}
	}
	if(lose==child){
		DP[u].y++;
	}
	else{
		DP[u].x++;
	}
	int sumx=0,sumy=0,sumz=0;
	for(int v:graph[u]){
		if(res[v]){
			lose--;
		}
		sumx+=DP[v].x,sumy+=DP[v].y,sumz+=DP[v].z;
		if(lose==child-1){
			DP[u].x+=(sz[v]-DP[v].x-DP[v].y-DP[v].z);
			DP[u].y+=DP[v].z;
			DP[u].z+=DP[v].y;
		}
		else{
			DP[u].x+=sz[v];
		}
		if(res[v]){
			lose++;
		}
	}
	x+=DP[u].x,y+=DP[u].y,z+=DP[u].z;
	x%=mod,y%=mod,z%=mod;
	if(u==1){
		x1=DP[u].x,y1=DP[u].y,z1=DP[u].z;
	}
	for(int i=0;i<graph[u].size();i++){
		int v=graph[u][i];
		if(v==p){
			continue;
		}
		bruh oriu=DP[u],oriv=DP[v];
		bool resu=res[u],resv=res[v];
		int szu=sz[u],szv=sz[v];
		sz[u]-=sz[v];
		if(res[v]){
			lose--;
		}
		child--;
		sumx-=DP[v].x,sumy-=DP[v].y,sumz-=DP[v].z;
		if(lose==child){
			res[u]=false;
		}
		else{
			res[u]=true;
		}
		DP[u].x=0,DP[u].y=0,DP[u].z=0;
		if(lose==child){
			DP[u].x=(sz[u]-1-sumx-sumy-sumz);
			DP[u].y=sumz+1;
			DP[u].z=sumy;
		}
		else if(lose==child-1){
			int vc=-1;
			for(int j=0;j<candi.size();j++){
				if(candi[j]!=v){
					vc=candi[j];
					break;
				}
			}
			DP[u].x=sz[u]-sz[vc]+(sz[vc]-DP[vc].x-DP[vc].y-DP[vc].z);
			DP[u].y=DP[vc].z;
			DP[u].z=DP[vc].y;
		}
		else{
			DP[u].x=sz[u];
		}
		DFS(v,u);
		DP[u]=oriu,DP[v]=oriv;
		res[u]=resu,res[v]=resv;
		sz[u]=szu,sz[v]=szv;
		if(res[v]){
			lose++;
		}
		child++;
		sumx+=DP[v].x,sumy+=DP[v].y,sumz+=DP[v].z;
	}
}
struct matrix{
	int n,m;
	int mat[2][2];
};
matrix matmul(matrix a,matrix b){
	matrix c;
	memset(c.mat,0,sizeof(c.mat));
	c.n=a.n,c.m=b.m;
	for(int i=0;i<c.n;i++){
		for(int k=0;k<a.m;k++){
			for(int j=0;j<c.m;j++){
				c.mat[i][j]+=(1LL*a.mat[i][k]*b.mat[k][j]%mod);
				c.mat[i][j]%=mod;
			}
		}
	}
	return c;
}
matrix binpow(matrix n,long long k){
	matrix curr=n;
	for(int i=0;i<curr.n;i++){
		for(int j=0;j<curr.m;j++){
			if(i!=j){
				curr.mat[i][j]=0;
			}
			else{
				curr.mat[i][j]=1;
			}
		}
	}
	for(;k;k>>=1){
		if(k&1){
			curr=matmul(curr,n);
		}
		n=matmul(n,n);
	}
	return curr;
}
signed main(){
	ios_base::sync_with_stdio(false);
	cin.tie(NULL);
	long long n,d;
	cin >> n >> d;
	for(int i=0;i<n-1;i++){
		int u,v;
		cin >> u >> v;
		graph[u].push_back(v),graph[v].push_back(u);
	}
	if(n==2){
		cout << binpow(1LL*n*n%mod,d);
		return 0;
	}
	setup(1,-1);
	DFS(1,-1);
	matrix curr;
	matrix trans;
	trans.n=2,trans.m=2;
	trans.mat[0][0]=(1LL*n*n%mod),trans.mat[0][1]=(x+y)%mod;
	trans.mat[1][0]=0,trans.mat[1][1]=(z-y+mod)%mod;
	curr.n=1,curr.m=2;
	curr.mat[0][0]=n,curr.mat[0][1]=wcnt;
	trans=binpow(trans,d-1);
	curr=matmul(curr,trans);
	wcnt=curr.mat[0][1];
	x=x1,y=y1,z=z1;
	int lcnt=((1LL*binpow(1LL*n*n%mod,d-1)*n%mod)-wcnt+mod)%mod;
	int nxtw=0;
	nxtw+=(1LL*x*binpow(1LL*n*n%mod,d-1)%mod*n%mod);
	nxtw%=mod;
	nxtw+=(1LL*y*lcnt%mod);
	nxtw%=mod;
	nxtw+=(1LL*z*wcnt%mod);
	nxtw%=mod;
	wcnt=nxtw;
	cout << wcnt;
}

• Subtask 1: 0 / 0
• Subtask 2: 7 / 7
• Subtask 3: 8 / 8
• Subtask 4: 15 / 15
• Subtask 5: 15 / 15
• Subtask 6: 20 / 20
• Subtask 7: 20 / 20
• Subtask 8: 15 / 15