centroid.cpp

#include <iostream>
#include <array>
#include <vector>
#include <string>
#include <algorithm>
#include <random>
#include <utility>
#include <tuple>
#include <chrono>

using namespace std;

using ll = long long;

int const maxn = 202020;

vector<int> g[maxn];
int tam[maxn], vis[maxn];

void dfs_tam(int u, int p = -1) {
	tam[u] = 1;
	for (int v : g[u]) {
		if (v == p or vis[v]) continue;
		dfs_tam(v, u);
		tam[u] += tam[v];
	}
}

int dfs_cent(int u, int tot, int p = -1) {
	for (int v : g[u]) {
		if (v == p or vis[v]) continue;
		if (tam[v] > tot / 2) return dfs_cent(v, tot, u);
	}
	return u;
}

int centroid(int u) {
	dfs_tam(u);
	return dfs_cent(u, tam[u]);
}

int d_cnt[maxn];

ll dfs_get(int u, int d_at, int d, int p = -1) {
	if (d < d_at) return 0;
	ll ans = d_cnt[d - d_at];
	for (int v : g[u]) {
		if (v == p or vis[v]) continue;
		ans += dfs_get(v, d_at, d, u);
	}
	return ans;
}

void dfs_add(int u, int d_at, int to_add, int p = -1) {
	d_cnt[d_at] += to_add;
	for (int v : g[u]) {
		if (v == p or vis[v]) continue;
		dfs_add(v, d_at+1, to_add, u);
	}
}

ll solve(int u, int d) {
	u = centroid(u);
	vis[u] = 1;

	ll ans = 0;
	d_cnt[0] = 1;
	for (int v : g[u]) {
		if (vis[v]) continue;

		ans += dfs_get(v, 1, d);
		dfs_add(v, 1, 1);
	}
	dfs_add(u, 0, -1);

	for (int v : g[u]) {
		if (vis[v]) continue;

		ans += solve(v, d);
	}

	return ans;
}

int main() {
	// calcula o numero de pares de vertices (u, v) tal que a distancia entre u e v seja d

	int d;
	// le entrada
	
	cout << solve(0, d)  << "\n";
}

centroid_query.cpp

#include <iostream>
#include <array>
#include <vector>
#include <string>
#include <algorithm>
#include <random>
#include <utility>
#include <tuple>
#include <chrono>

using namespace std;

using ll = long long;

int const maxn = 202020;

vector<int> g[maxn];
int tam[maxn], vis[maxn];

void dfs_tam(int u, int p = -1) {
	tam[u] = 1;
	for (int v : g[u]) {
		if (v == p or vis[v]) continue;
		dfs_tam(v, u);
		tam[u] += tam[v];
	}
}

int dfs_cent(int u, int tot, int p = -1) {
	for (int v : g[u]) {
		if (v == p or vis[v]) continue;
		if (tam[v] > tot / 2) return dfs_cent(v, tot, u);
	}
	return u;
}

int centroid(int u) {
	dfs_tam(u);
	return dfs_cent(u, tam[u]);
}

int anc[maxn];

void build_tree(int u, int p = -1) {
	u = centroid(u);
	vis[u] = 1;
	anc[u] = p;

	for (int v : g[u]) {
		if (vis[v]) continue;
		build_tree(v, u);
	}
}

int dist(int u, int v) {
	// retorna distancia de u para v
}

int min_dist[maxn];

void upd(int u) {
	int u_orig = u;
	while (u != -1) {
		min_dist[u] = min(min_dist[u], dist(u, u_orig));
		u = anc[u];
	}
}

int query(int u) {
	int ans = 0x3f3f3f3f, u_orig = u;
	while (u != -1) {
		ans = min(ans, dist(u_orig, u) + min_dist[u]);
		u = anc[u];
	}
	return ans;
}


int main() {
	// dado arvore com todos os vertices desligados, responder as seguintes queries
	// upd(u): liga o vertice u, faz nada caso ja esteja ligado
	// query(u): encontra a distancia minima de algum vertice ligado para o vertice u


	// le entrada

	build_tree(0);
}

dfs_reroot.cpp

#include <iostream>
#include <array>
#include <vector>
#include <string>
#include <algorithm>
#include <random>
#include <utility>
#include <tuple>
#include <chrono>

using namespace std;

using ll = long long;

int const maxn = 202020;

vector<int> g[maxn];
ll tam[maxn], dp[maxn], ans[maxn];

void dp_init(int u, int p = -1) {
	tam[u] = 1;
	dp[u] = 0;

	for (int v : g[u]) {
		if (v == p) continue;

		dp_init(v, u);
		tam[u] += tam[v];
		dp[u] += dp[v] + tam[v];
	}
}

void reroot(int old_root, int new_root) {
	// O(1)
	dp[old_root] -= dp[new_root] + tam[new_root];
	tam[old_root] -= tam[new_root];
	tam[new_root] += tam[old_root];
	dp[new_root] += dp[old_root] + tam[old_root];
}

void dfs(int u, int p = -1) {
	ans[u] = dp[u];

	for (int v : g[u]) {
		if (v == p) continue;

		reroot(u, v);
		dfs(v, u);
		reroot(v, u);
	}
}

int main() {
	ios_base::sync_with_stdio(false), cin.tie(nullptr);
	
	// calcular pra cada vertice u a soma dist(u, v) pra cada vertice v
	
	// ler entrada

	dp_init(0);
	dfs(0);

	// ans[u] = soma dos dist[u, v] pra todos os vertices v
}

dfs_reroot_2.cpp

#include <iostream>
#include <array>
#include <vector>
#include <string>
#include <algorithm>
#include <random>
#include <utility>
#include <tuple>
#include <chrono>

#define ff first
#define ss second

#define pb push_back
#define eb emplace_back
#define clr(x, c) memset((x), (c), sizeof((x)))

using namespace std;

template<class T> void DBG(T&& x) { cerr << x << " "; }
template<class T, class...Args> void DBG(T&& x, Args&&... args) { DBG(x); DBG(args...); }
#define DBG(...) do {cerr << "[" << #__VA_ARGS__ << "]: "; DBG(__VA_ARGS__); cerr << endl; } while (0) 

#ifdef _WIN32
#define getchar_unlocked _getchar_nolock
#endif

template<class num> inline void rd(num& x) {
	char c, neg = 0; while(isspace(c = getchar_unlocked()));
	if(!isdigit(c)) neg = (c == '-'), x = 0;
	else x = c - '0';
	while(isdigit(c = getchar_unlocked())) x = (x << 3) + (x << 1) + c - '0';
	x = neg ? -x : x; }
template <class Ty, class... Args> inline void rd(Ty& x, Args&... args) { rd(x); rd(args...); }


using ll = long long;
using ii = pair<int, int>;

using Graph = vector<vector<int>>;
using WGraph = vector<vector<ii>>;

template<typename T = int>
constexpr T inf = 0x3f3f3f3f;

template<>
constexpr ll inf<ll> = 0x3f3f3f3f3f3f3f3f;

seed_seq seq {
    (uint64_t) chrono::duration_cast<chrono::nanoseconds>(
    	chrono::high_resolution_clock::now().
    	time_since_epoch()).count(),
    (uint64_t) random_device{}(),
    (uint64_t) 17
};

mt19937 rng{seq};
template<class Ty> Ty randint(Ty a, Ty b) { return uniform_int_distribution<Ty>(a, b)(rng); }

template<class... Ts> struct overloaded : Ts... { using Ts::operator()...; };
template<class... Ts> overloaded(Ts...) -> overloaded<Ts...>;

template<class F>
struct y_combinator_result{
	F f;
	template<class T> explicit y_combinator_result(T &&f): f(forward<T>(f)){ }
	template<class ...Args> decltype(auto) operator()(Args &&...args){ return f(ref(*this), forward<Args>(args)...); }
};

template<class F>
decltype(auto) y_combinator(F &&f){
	return y_combinator_result<decay_t<F>>(forward<F>(f));
}

int const maxn = 202020;

vector<int> g[maxn], pref[maxn], suf[maxn];
ll dp[maxn], a[maxn], ans[maxn];

void dp_init(int u, int p = -1) {
	dp[u] = a[u];

	for (int v : g[u]) {
		if (v == p) continue;

		dp_init(v, u);
		dp[u] = gcd(dp[u], dp[v]);
	}

	pref[u].resize(g[u].size());
	suf[u].resize(g[u].size());

	for (int i = 1; i < (int)g[u].size(); i++) {
		pref[u][i] = gcd(pref[u][i-1], (g[u][i-1] == p ? 0 : a[g[u][i-1]]));
	}
	for (int i = (int)g[u].size() - 2; i >= 0 ; i--) {
		suf[u][i] = gcd(suf[u][i+1], (g[u][i+1] == p ? 0 : a[g[u][i+1]]));
	}

	// valor do gcd da subarvore do vertice u, sem subarvore i = gcd(pref[u][i], suf[u][i], a[u]);
}

void dfs(int u, int p = -1, int dp_up = 0) {
	ans[u] = dp_up;
	// armazenar o valor de dp[u] em um vetor persistente

	for (int i = 0; i < (int)g[u].size(); i++) {
		int v = g[u][i];
		if (v == p) continue;

		ans[u] += dp[v];
		// atualizar valor passado
		int new_up = gcd(dp_up, gcd(pref[u][i], gcd(suf[u][i], a[i])));
		dfs(v, u, new_up);
	}
}

int main() {
	ios_base::sync_with_stdio(false), cin.tie(nullptr);
	
	// calcular pra cada vertice u a soma dos gcds de cada arvore criada apos
	// remocao do vertice u
	
	// ler entrada

	dp_init(0);
	dfs(0);
}