Submission #1366631

#TimeUsernameProblemLanguageResultExecution timeMemory
1366631gelastropodJOI Tour 2 (JOI26_joitour)C++20
100 / 100
3571 ms138916 KiB
#pragma GCC optimize("O3,inline")
#include <bits/stdc++.h>
using namespace std;
#define int long long

int N, M, Q, cnt = 0, sqrtN;
vector<int> A, B, d, tour, pre, post, gguys, bigW, gcj, itself;
vector<vector<int>> adjlist, p2;
vector<pair<pair<int, int>, int>> path;
vector<vector<pair<pair<int, int>, int>>> fguys;
vector<pair<int, int>> touridxs;
pair<int, int> endpts = { 0, 0 };

#define f(a, b, c, d) if (a != -1) fguys[a].push_back({ { b, c }, d });
#define g(a, b, c, d) for (int i = a; i != b; i = p2[0][i]) gguys[i] += c, itself[2 * A[i]] += c * d;

void dfs(int n, int p) {
	pre[n] = cnt++;
	tour.push_back(n);
	for (int i : adjlist[n]) {
		if (i == p) continue;
		p2[0][i] = n;
		d[i] = d[n] + 1;
		dfs(i, n);
	}
	post[n] = cnt++;
	tour.push_back(n);
}

int lca(int a, int b) {
	if (d[a] < d[b]) swap(a, b);
	for (int i = 0; i < 20; i++) if ((d[a] - d[b]) & (1LL << i)) a = p2[i][a];
	if (a == b) return a;
	for (int i = 19; i >= 0; i--) if (p2[i][a] != p2[i][b]) a = p2[i][a], b = p2[i][b];
	return p2[0][a];
}

bool comp(pair<pair<int, int>, int> a, pair<pair<int, int>, int> b) {
	auto &ia = touridxs[a.second], &ib = touridxs[b.second];
	if (ia.first / sqrtN == ib.first / sqrtN) return ia.second < ib.second;
	return ia.first < ib.first;
}

void movepath(int v, int f, int w) {
	if (f) swap(endpts.first, endpts.second);
	auto [s, t] = endpts;
	int x = lca(s, t), y = lca(t, v), z = lca(s, v);
	if (x == z) {
		f(s, t, y, -w); f(x, t, y, w); f(p2[0][x], t, y, w); g(t, y, -w, 0);
		f(s, v, y, w); f(x, v, y, -w); f(p2[0][x], v, y, -w); g(v, y, w, 0);
	}
	else {
		f(s, t, x, -w); f(x, t, x, w); f(p2[0][x], t, x, w); g(t, x, -w, 0);
		if (y == z) {
			f(s, p2[0][x], p2[0][y], w); g(p2[0][x], p2[0][y], -w, 1);
			f(s, v, y, w); f(y, v, y, -w); f(p2[0][y], v, y, -w); g(v, y, w, 0);
		}
		else {
			f(s, p2[0][z], p2[0][x], -w); g(p2[0][z], p2[0][x], w, 1);
			f(s, v, z, w); f(z, v, z, -w); f(p2[0][z], v, z, -w); g(v, z, w, 0);
		}
	}
	endpts.second = v;
	if (f) swap(endpts.first, endpts.second);
}

void mk(int n, int p) {
	for (int i = 0; i < Q; i++) if (0 <= B[i] - A[n] && B[i] - A[n] <= N) bigW[i] -= gcj[B[i] - A[n]];
	for (auto& i : fguys[n]) for (; i.first.first != i.first.second; i.first.first = p2[0][i.first.first]) gcj[A[i.first.first]] += i.second;
	for (int i : adjlist[n]) {
		if (i == p) continue;
		gcj[A[i]] += gguys[i];
		mk(i, n);
	}
	for (int i = 0; i < Q; i++) if (0 <= B[i] - A[n] && B[i] - A[n] <= N) bigW[i] += gcj[B[i] - A[n]];
}

signed main() {
	ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);
	int a, b; cin >> N;
	sqrtN = sqrt(N);
	adjlist.resize(N, vector<int>());
	p2.resize(20, vector<int>(N, -1));
	pre.resize(N, -1);
	post.resize(N, -1);
	d.resize(N, 0);
	fguys.resize(N, vector<pair<pair<int, int>, int>>());
	gguys.resize(N, 0);
	gcj.resize(N + 1, 0);
	itself.resize(2 * N + 1, 0);
	for (int i = 0; i < N; i++) {
		cin >> a;
		A.push_back(a);
	}
	for (int i = 0; i < N - 1; i++) {
		cin >> a >> b; a--, b--;
		adjlist[a].push_back(b);
		adjlist[b].push_back(a);
	}
	dfs(0, -1);
	for (int i = 1; i < 20; i++) {
		for (int j = 0; j < N; j++) {
			if (p2[i - 1][j] == -1) continue;
			p2[i][j] = p2[i - 1][p2[i - 1][j]];
		}
	}
	cin >> M;
	for (int i = 0; i < M; i++) {
		cin >> a >> b; a--, b--;
		if (pre[a] > pre[b]) swap(a, b);
		int clca = lca(a, b);
		if (a == clca) touridxs.push_back({ pre[a], pre[b] });
		else touridxs.push_back({ post[a], pre[b] });
		path.push_back({ { a, b }, i });
	}
	sort(path.begin(), path.end(), comp);
	for (int i = 0; i < M; i++) {
		movepath(path[i].first.first, 1, M - i);
		movepath(path[i].first.second, 0, M - i);
	}
	cin >> Q;
	bigW.resize(Q, 0);
	for (int i = 0; i < Q; i++) {
		cin >> a;
		B.push_back(a);
	}
	mk(0, -1);
	for (int i = 0; i < Q; i++) bigW[i] += itself[B[i]];
	for (int i : bigW) cout << i << '\n';
}
#Result Execution timeMemoryGrader output
Fetching results...
#Result Execution timeMemoryGrader output
Fetching results...
#Result Execution timeMemoryGrader output
Fetching results...
#Result Execution timeMemoryGrader output
Fetching results...
#Result Execution timeMemoryGrader output
Fetching results...
#Result Execution timeMemoryGrader output
Fetching results...
#Result Execution timeMemoryGrader output
Fetching results...
#Result Execution timeMemoryGrader output
Fetching results...
#Result Execution timeMemoryGrader output
Fetching results...
#Result Execution timeMemoryGrader output
Fetching results...
#Result Execution timeMemoryGrader output
Fetching results...
#Result Execution timeMemoryGrader output
Fetching results...