Submission #292927

#	Time	Username	Problem	Language	Result	Execution time	Memory
292927		arnold518	Star Trek (CEOI20_startrek)	C++14	100 / 100	138 ms	16788 KiB

startrek

#include <bits/stdc++.h>
using namespace std;

typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;

const int MAXN = 1e5;
const ll MOD = 1e9+7;

int N; ll D;
vector<int> adj[MAXN+10];

typedef vector<vector<ll>> Mat;

Mat operator * (const Mat &p, const Mat &q)
{
	Mat ret=vector<vector<ll>>(p.size(), vector<ll>(q[0].size()));
	assert(p[0].size()==q.size());
	for(int i=0; i<p.size(); i++) for(int j=0; j<q[0].size(); j++)
	{
		for(int k=0; k<p[0].size(); k++)
		{
			ret[i][j]+=p[i][k]*q[k][j]%MOD;
			ret[i][j]%=MOD;
		}
	}
	return ret;
}

Mat mypow(Mat x, ll y)
{
	if(y==0) return {{1, 0}, {0, 1}};
	if(y%2) return mypow(x, y-1)*x;
	Mat t=mypow(x, y/2);
	return t*t;
}

ll ans;

ll a, b, c, d, p, q, x, y;

int dp[MAXN+10], dp2[MAXN+10], zero;
int cnt0[MAXN+10], cnt1[MAXN+10];

void dfs(int now, int bef)
{
	for(int nxt : adj[now])
	{
		if(nxt==bef) continue;
		dfs(nxt, now);
		dp[now]+=!dp[nxt];
	}
	zero+=!dp[now];

	dp2[now]=!dp[now];
	for(int nxt : adj[now])
	{
		if(nxt==bef) continue;
		if(dp[nxt]==0) cnt0[now]+=dp2[nxt];
		else if(dp[nxt]==1) cnt1[now]+=dp2[nxt];
	}
	if(dp[now]==0) dp2[now]+=cnt1[now];
	else if(dp[now]==1) dp2[now]+=cnt0[now];
}

void dfs2(int now, int bef)
{
	//printf("%d : dp %d dp2 %d zero %d\n", now, dp[now], dp2[now], zero);
	//for(int i=1; i<=N; i++) printf("%d : %d / ", i, dp2[i]); printf("\n");
	if(dp[now]==0) c++;
	else a++;

	if(dp[now]==0) d+=N-zero;
	else b+=N-zero;

	if(dp[now]==0)
	{
		b+=dp2[now];
		d+=zero-dp2[now];
	}
	else
	{
		d+=dp2[now];
		b+=zero-dp2[now];
	}

	//printf("%lld %lld %lld %lld\n", a, b, c, d);
	if(now==1) { x=a; y=b; }

	for(int nxt : adj[now])
	{
		if(nxt==bef) continue;

		int m1=dp2[now], m2=dp2[nxt];
		int m3=cnt1[now], m4=cnt0[now], m5=cnt1[nxt], m6=cnt0[nxt];

		zero-=!dp[now];
		zero-=!dp[nxt];

		dp[now]-=!dp[nxt];
		if(dp[nxt]==0) cnt0[now]-=dp2[nxt];
		else if(dp[nxt]==1) cnt1[now]-=dp2[nxt];

		dp2[now]=!dp[now];
		if(dp[now]==0) dp2[now]+=cnt1[now];
		else if(dp[now]==1) dp2[now]+=cnt0[now];

		dp[nxt]+=!dp[now];
		if(dp[now]==0) cnt0[nxt]+=dp2[now];
		else if(dp[now]==1) cnt1[nxt]+=dp2[now];

		dp2[nxt]=!dp[nxt];
		if(dp[nxt]==0) dp2[nxt]+=cnt1[nxt];
		else if(dp[nxt]==1) dp2[nxt]+=cnt0[nxt];

		zero+=!dp[nxt];
		zero+=!dp[now];

		dfs2(nxt, now);

		zero-=!dp[now];
		zero-=!dp[nxt];

		dp[nxt]-=!dp[now];
		dp[now]+=!dp[nxt];
		zero+=!dp[nxt];
		zero+=!dp[now];

		dp2[now]=m1;
		dp2[nxt]=m2;

		cnt1[now]=m3; cnt0[now]=m4;
		cnt1[nxt]=m5; cnt0[nxt]=m6;
	}
}

int main()
{
	scanf("%d%lld", &N, &D);

	for(int i=1; i<N; i++)
	{
		int u, v;
		scanf("%d%d", &u, &v);
		adj[u].push_back(v);
		adj[v].push_back(u);
	}

	dfs(1, 1);
	dfs2(1, 1);

	p=a; q=c;
	a*=N; c*=N; x*=N;

	a%=MOD; b%=MOD; c%=MOD; d%=MOD; x%=MOD; y%=MOD; p%=MOD; q%=MOD;

	Mat M2={{a, b}, {c, d}};
	Mat M1={{x, y}, {0, 0}};
	Mat M3={{p}, {q}};

	Mat MM=M1*mypow(M2, D-1)*M3;

	//printf("!a %lld b %lld c %lld d %lld / x %lld y %lld\n", a, b, c, d, x, y);

	printf("%lld\n", MM[0][0]);
}

Compilation message (stderr)

startrek.cpp: In function 'Mat operator*(const Mat&, const Mat&)':
startrek.cpp:20:16: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::vector<long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   20 |  for(int i=0; i<p.size(); i++) for(int j=0; j<q[0].size(); j++)
      |               ~^~~~~~~~~
startrek.cpp:20:46: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   20 |  for(int i=0; i<p.size(); i++) for(int j=0; j<q[0].size(); j++)
      |                                             ~^~~~~~~~~~~~
startrek.cpp:22:17: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   22 |   for(int k=0; k<p[0].size(); k++)
      |                ~^~~~~~~~~~~~
startrek.cpp: In function 'int main()':
startrek.cpp:140:7: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
  140 |  scanf("%d%lld", &N, &D);
      |  ~~~~~^~~~~~~~~~~~~~~~~~
startrek.cpp:145:8: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
  145 |   scanf("%d%d", &u, &v);
      |   ~~~~~^~~~~~~~~~~~~~~~

• 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