Submission #91150

#	Time	Username	Problem	Language	Result	Execution time	Memory
91150		Just_Solve_The_Problem	Shymbulak (IZhO14_shymbulak)	C++11	100 / 100	796 ms	21996 KiB

shymbulak

#include <stdio.h>
#include <vector>
#include <utility>
#include <memory.h>
#include <iostream>
#include <map>
#include <queue>
#include <set>

using namespace std;

#define ok puts("ok");
#define ll long long

const int N = (int)2e5 + 7;

vector < int > gr[N];
int n;
int used[N];
ll ans = 0;

int diam;
vector < int > node;
int deg[N];
int sz[N], isdel[N];
map < int, ll > mp1, mp2;

void precalc(int v, int pr) {
	sz[v] = 1;
	for (int to : gr[v]) {
		if (to == pr || used[to] == 2 || isdel[to]) continue;
		precalc(to, v);
		sz[v] += sz[to];
	}
}

int findcentroid(int v) {
	int cur = v;
	int pr = v;
	bool run = 1;
	while (run) {
		run = 0;
		for (int to : gr[cur]) {
			if (to == pr || (used[to] == 2) || isdel[to]) continue;
			if (sz[to] * 2 >= sz[v]) {
				run = 1;
				pr = cur;
				cur = to;
				break;
			}
		}
	}
	return cur;
}

int cnt[N][20];

void calc(int v, int pr, int val, int d, int dep) {
	cnt[d][dep] += val;
	for (int to : gr[v]) {
		if (to == pr || used[to] == 2 || isdel[to]) continue;
		calc(to, v, val, d + 1, dep);
	}
} 

ll res;

void calc1(int v, int pr, int d, int dep) {
	res += cnt[diam - d][dep];
	for (int to : gr[v]) {
		if (to == pr || used[to] == 2 || isdel[to]) continue;
		calc1(to, v, d + 1, dep);
	}
}

void decompose(int v, int dep) {
	precalc(v, v);
	int cen = findcentroid(v);
	isdel[cen] = 1;
	calc(cen, cen, 1, 0, dep);
	res = 0;
	res += cnt[diam][dep];
	for (int to : gr[cen]) {
		if (isdel[to] || used[to] == 2) continue;
		calc(to, cen, -1, 1, dep);
		calc1(to, cen, 1, dep);
		calc(to, cen, 1, 1, dep);
	}
	res /= 2;
	ans += res;
	calc(cen, cen, -1, 0, dep);
	for (int to : gr[cen]) {
		if (isdel[to] || used[to] == 2) continue;
		decompose(to, dep + 1);
	}
}

pair < int, int > mxh[N];

void dfs(int v, int pr) {
	bool asd = 0;
	for (int to : gr[v]) {
		if (to == pr || used[to] == 2) continue;
		dfs(to, v);
		asd = 1;
		if (mxh[to].first == mxh[v].first) {
			mxh[v].second += mxh[to].second;
		} else if (mxh[to].first > mxh[v].first) {
			mxh[v] = mxh[to];
		}
	}
	if (!asd) {
		mxh[v] = {0, 1};
	} else
		mxh[v].first++;
}

int us[N];

void dfs1(int v, int pr) {
	node.push_back(v);
	us[v] = 1;
	for (int to : gr[v]) {
		if (us[to] || used[to] != 2) continue;
		dfs1(to, v);
	}
}

int dd[N];
pair < int, int > tamer;

void dfs3(int v, int pr) {
	if (dd[v] > tamer.first) {
		tamer.first = dd[v];
		tamer.second = v;
	}
	for (int to : gr[v]) {
		if (to == pr || used[to] == 2) continue;
		dd[to] = dd[v] + 1;
		dfs3(to, v);
	}
}

main() {
	//freopen("shymbulak.in", "r", stdin);
	//freopen("shymbulak.out", "w", stdout);
	scanf("%d", &n);
	for (int i = 1; i <= n; i++) {
		int u, v;
		scanf("%d %d", &u, &v);
		gr[u].push_back(v);
		gr[v].push_back(u);
		deg[u]++;
		deg[v]++;
	}
	for (int i = 1; i <= n; i++) {
		if (gr[i].size() == 1) {
			node.push_back(i);
		}
	}
	memset(used, 0, sizeof used);
	for (int i = 0; i < n; i++) {
		if (i == node.size()) {
			break;
		}
		int v = node[i];
		used[v] = 1;
		for (int to : gr[v]) {
			if (!used[to]) {
				deg[to]--;
				if (deg[to] <= 1) {
					node.push_back(to);
				}
			}
		}
	}
	for (int i = 1; i <= n; i++) {
		if (!used[i]) {
			used[i] = 2;
		}
	}
	while (!node.empty()) node.pop_back();
	for (int i = 1; i <= n; i++) {
		if (used[i] == 2) {
			dfs1(i, i);
			break;
		}
	}
	for (int i = 1; i <= n; i++) {
		if (used[i] == 2) {
			dfs(i, i);
		}
	}
	////
	int h1, cur;
	multiset < pair < int, int > > S, s1;
	h1 = node.size() / 2;
	for (int i = 0; i < h1; i++) {
		S.insert({mxh[node[i]].first - i, node[i]});
	}
	for (int i = h1; i < node.size(); i++) {
		pair < int, int > v = *S.rbegin();
		diam = max(diam, mxh[node[i]].first + i + v.first);
		int lst = i - h1;
		S.erase(S.find({mxh[node[lst]].first - lst, node[lst]}));
		S.insert({mxh[node[i]].first - i, node[i]});
	}
	for (int i = node.size() - h1; i < node.size(); i++) {
		s1.insert({(int)node.size() - i + mxh[node[i]].first, node[i]});
	}
	for (int i = 0; i < h1; i++) {
		pair < int, int > v = *s1.rbegin();
		diam = max(diam, mxh[node[i]].first + i + v.first);
		int lst = i - h1 + (int)node.size();
		s1.erase(s1.find({(int)node.size() - lst + mxh[node[lst]].first, node[lst]}));
		s1.insert({mxh[node[i]].first - i, node[i]});
	}
	// finding tree diameter
	for (int i = 1; i <= n; i++) {
		if (used[i] == 2) {
			tamer.first = -1;
			used[i] = 1;
			dd[i] = 0;
			dfs3(i, i);
			//cout << i << ' ' << tamer.first << ' ' << tamer.second << endl;
			tamer.first = -1;
			dd[tamer.second] = 0;
			dfs3(tamer.second, tamer.second);
			//cout << i << ' ' << tamer.first << ' ' << tamer.second << endl;
			diam = max(diam, tamer.first);
			used[i] = 2;
		}
	}
	for (int i = 1; i <= n; i++) {
		if (used[i] == 2) {
			used[i] = 1;
			decompose(i, 0);
			used[i] = 2;
		}
	}
	h1 = node.size() / 2;
	if (node.size() & 1 ^ 1) h1--;
	for (int i = 0; i < h1; i++) {
		mp1[-i + mxh[node[i]].first] += mxh[node[i]].second;
	}
	cur = h1;
	res = 0;
	while (cur < (int)node.size()) {
		res += mp1[diam - mxh[node[cur]].first - cur] * 1LL * mxh[node[cur]].second;
		int lst = cur - h1;
		mp1[-lst + mxh[node[lst]].first] -= mxh[node[lst]].second;
		mp1[-cur + mxh[node[cur]].first] += mxh[node[cur]].second;
		cur++;
	}
	for (int i = (int)node.size() - h1; i < node.size(); i++) {
		mp2[(int)node.size() - i + mxh[node[i]].first] += mxh[node[i]].second;
	}
	cur = 0;
	while (cur < h1) {
		res += mp2[diam - mxh[node[cur]].first - cur] * 1LL * mxh[node[cur]].second;
		int lst = cur - h1 + (int)node.size();
		mp2[(int)node.size() - lst + mxh[node[lst]].first] -= mxh[node[lst]].second;
		mp2[-cur + mxh[node[cur]].first] += mxh[node[cur]].second;
		cur++;
	}
	// assert(res & 1 ^ 1);
	ans += res;
	res = 0;
	if ((int)node.size() & 1 ^ 1) {
		h1 = node.size() / 2;
		for (int i = 0; i < node.size(); i++) {
			int lst = i - h1;
			if (lst < 0) lst += (int)node.size();
			if (diam - h1 == mxh[node[i]].first + mxh[node[lst]].first) {
				res += mxh[node[i]].second * 1LL * mxh[node[lst]].second;
			}
		}
		ans += res;
	}
	cout << ans;
}

Compilation message (stderr)

shymbulak.cpp:144:6: warning: ISO C++ forbids declaration of 'main' with no type [-Wreturn-type]
 main() {
      ^
shymbulak.cpp: In function 'int main()':
shymbulak.cpp:163:9: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
   if (i == node.size()) {
       ~~^~~~~~~~~~~~~~
shymbulak.cpp:201:21: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
  for (int i = h1; i < node.size(); i++) {
                   ~~^~~~~~~~~~~~~
shymbulak.cpp:208:35: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
  for (int i = node.size() - h1; i < node.size(); i++) {
                                 ~~^~~~~~~~~~~~~
shymbulak.cpp:242:18: warning: suggest parentheses around arithmetic in operand of '^' [-Wparentheses]
  if (node.size() & 1 ^ 1) h1--;
      ~~~~~~~~~~~~^~~
shymbulak.cpp:255:40: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
  for (int i = (int)node.size() - h1; i < node.size(); i++) {
                                      ~~^~~~~~~~~~~~~
shymbulak.cpp:269:23: warning: suggest parentheses around arithmetic in operand of '^' [-Wparentheses]
  if ((int)node.size() & 1 ^ 1) {
      ~~~~~~~~~~~~~~~~~^~~
shymbulak.cpp:271:21: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
   for (int i = 0; i < node.size(); i++) {
                   ~~^~~~~~~~~~~~~
shymbulak.cpp:147:7: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
  scanf("%d", &n);
  ~~~~~^~~~~~~~~~
shymbulak.cpp:150:8: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
   scanf("%d %d", &u, &v);
   ~~~~~^~~~~~~~~~~~~~~~~

• Subtask 1: 30 / 30
• Subtask 2: 20 / 20
• Subtask 3: 50 / 50