Submission #693069

#	Time	Username	Problem	Language	Result	Execution time	Memory
693069		US3RN4M3	Wells (CEOI21_wells)	C++17	0 / 100	19 ms	36948 KiB

wells

#include<bits/stdc++.h>
using namespace std;
using ll = long long;
const ll mod = 1000000007;
const int mx = 1500005;
bool valid[mx];
vector<int> graph[mx];
bool vis[mx];
int par[mx];
int dist[mx];
bool deg3[mx];
int n, k;
void dfs1(int cur, int par) {
	for(int nxt : graph[cur]) if(nxt != par) {
		dist[nxt] = dist[cur] + 1;
		dfs1(nxt, cur);
	}
	for(int nxt : graph[cur]) if(nxt != par) {
		dist[cur] = max(dist[cur], dist[nxt]);
	}
}
void dfs2(int cur, int par) {
	for(int nxt : graph[cur]) if(nxt != par) {
		dist[nxt] = dist[cur] + 1;
		dfs2(nxt, cur);
	}
}
main() {
	cin >> n >> k;
	for(int i = 0; i < n - 1; i++) {
		int s, d; cin >> s >> d;
		graph[s].push_back(d);
		graph[d].push_back(s);
	}
	vis[1] = true;
	vector<int> q{1};
	int qidx = 0;
	while(qidx < q.size()) {
		int cur = q[qidx++];
		for(int nxt : graph[cur]) {
			if(!vis[nxt]) {
				q.push_back(nxt);
				vis[nxt] = true;
			}
		}
	}
	int far = q.back();
	memset(vis, false, sizeof(vis));
	vis[far] = true;
	q = {far};
	qidx = 0;
	while(qidx < q.size()) {
		int cur = q[qidx++];
		for(int nxt : graph[cur]) {
			if(!vis[nxt]) {
				vis[nxt] = true;
				par[nxt] = cur;
				dist[nxt] = dist[cur] + 1;
				q.push_back(nxt);
			}
		}
	}
	int far2 = q.back();
	int width = dist[far2] + 1;
	int rad;
	if(width % 2 == 0) {
		rad = width / 2;
		int centre1 = far2;
		for(int i = 0; i < rad - 1; i++) {
			centre1 = par[centre1];
		}
		int centre2 = par[centre1];
		dist[centre1] = 0;
		dist[centre2] = 0;
		dfs1(centre1, centre2);
		dfs1(centre2, centre1);
		rad++;
	} else {
		rad = width / 2;
		int centre = far2;
		for(int i = 0; i < rad; i++) {
			centre = par[centre];
		}
		dist[centre] = 0;
		dfs1(centre, 0);
	}
	ll mul = 1;
	for(int i = 1; i <= n; i++) {
		if(dist[i] + rad + 1 < k) {
			mul += mul;
			if(mul >= mod) mul -= mod;
		} else {
			valid[i] = true;
		}
	}
	int root = -1;
	for(int i = 1; i <= n; i++) {
		if(valid[i]) {
			root = i;
		}
	}
	if(root == -1) {
		cout << "NO" << endl;
		cout << 0 << endl;
		return 0;
	}
	bool is_line = true;
	for(int i = 1; i <= n; i++) if(valid[i]) {
		int deg = 0;
		for(int nxt : graph[i]) if(valid[nxt]) deg++;
		if(deg >= 3) {
			deg3[i] = true;
			is_line = false;
			root = i;
		}
	}
	if(is_line) {
		cout << "YES" << endl;
		ll ans = (mul * k) % mod;
		cout << ans << endl;
		return 0;
	}
	dist[root] = 0;
	dfs2(root, 0);
	for(int i = 1; i <= n; i++) if(valid[i] && deg3[i]) {
		if(k & 1) {
			if(dist[i] % k != 0) {
				cout << "NO" << endl << 0 << endl;
				return 0;
			}
		} else {
			if(dist[i] % (k / 2) != 0) {
				cout << "NO" << endl << 0 << endl;
				return 0;
			}
		}
	}
	mul += mul;
	mul %= mod;
	cout << "YES" << endl << mul << endl;
}

Compilation message (stderr)

wells.cpp:28:1: warning: ISO C++ forbids declaration of 'main' with no type [-Wreturn-type]
   28 | main() {
      | ^~~~
wells.cpp: In function 'int main()':
wells.cpp:38:13: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   38 |  while(qidx < q.size()) {
      |        ~~~~~^~~~~~~~~~
wells.cpp:52:13: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   52 |  while(qidx < q.size()) {
      |        ~~~~~^~~~~~~~~~

• Subtask 1: 0 / 30
• Subtask 2: 0 / 20
• Subtask 3: 0 / 20
• Subtask 4: 0 / 30