Submission #660224

#	Time	Username	Problem	Language	Result	Execution time	Memory
660224		Sohsoh84	Wells (CEOI21_wells)	C++17	0 / 100	16 ms	23892 KiB

wells

#include <bits/stdc++.h>

using namespace std;

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

#define all(x)			(x).begin(),(x).end()
#define X			first
#define Y			second
#define sep			' '
#define endl			'\n'
#define debug(x)		cerr << #x << ": " <<  x << endl;

const ll MAXN = 1e6 + 10;
const ll MOD = 1e9 + 7;

inline ll poww(ll a, ll b) {
	ll ans = 1;
	while (b) {
		if (b & 1) ans = ans * a % MOD;
		b >>= 1;
		a = a * a % MOD;
	}

	return ans;
}

inline void push_max(pll& a, int b) {
	if (b > a.X) {
		a.Y = a.X;
		a.X = b;
	} else if (b > a.Y) a.Y = b;
}

inline void push_min(pll& a, int b) {
	if (b < a.X) {
		a.Y = a.X;
		a.X = b;
	} else if (b < a.Y) a.Y = b;
}

pll dp_down[MAXN], dp_max[MAXN], dp_min[MAXN];
int n, k, dp_up[MAXN], H[MAXN];
set<vector<int>> st;
vector<int> adj[MAXN], tvec;
bool f[MAXN], flag;

void dfs_down(int v, int p) {
	for (int u : adj[v]) {
		if (u == p) continue;
		dfs_down(u, v);
		push_max(dp_down[v], dp_down[u].X + 1);
	}
}

void dfs_up(int v, int p) {
	if (p) {
		dp_up[v] = max(dp_up[p] + 1, (dp_down[p].X == dp_down[v].X + 1 ? dp_down[p].Y : dp_down[p].X) + 1);
	}

	f[v] = (dp_down[v].X + max(dp_up[v], dp_down[v].Y) + 1 >= k);
	for (int u : adj[v])
		if (u != p)
			dfs_up(u, v);
}

void dfs(int v, int p) {
	if (p) H[v] = H[p] + 1;
	bool tmp = (H[v] % k > 0);

	dp_max[v] = {0, 0};
	dp_min[v] = {MAXN, MAXN};
	if (!tmp) {
		dp_max[v] = {-MAXN, -MAXN};
		dp_min[v] = {0, MAXN};
		tvec.push_back(v);
	}

	for (int u : adj[v]) {
		if (u == p || !f[u]) continue;
		dfs(u, v);

		if (tmp) push_max(dp_max[v], dp_max[u].X + 1);
		push_min(dp_min[v], dp_min[u].X + 1);
	}

	if (dp_max[v].X + dp_max[v].Y + 1 >= k) flag = false;
	if (dp_min[v].X + dp_min[v].Y + 1 <= k) flag = false;
}

int main() {
	ios_base::sync_with_stdio(false); cin.tie(0); cout.tie(0);
	cin >> n >> k;
	for (int i = 1; i < n; i++) {
		int u, v;
		cin >> u >> v;
		adj[u].push_back(v);
		adj[v].push_back(u);
	}

	dfs_down(1, 0);
	dfs_up(1, 0);
	
	ll tans = 0;
	for (int i = 1; i <= n; i++) {
		if (!f[i]) {
			tans++;
			continue;
		}

		H[i] = 0;
		flag = true;
		tvec.clear();
		dfs(i, 0);

		if (flag) {
			sort(all(tvec));
			st.insert(tvec);
		}
	}

	if (tans == n) return cout << "YES" << endl << poww(2, tans) << endl, 0;
	cout << (st.empty() ? "NO" : "YES") << endl;
	cout << ll(st.size()) * poww(2, tans) % MOD << endl;
	return 0;
}

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