Submission #1369692

#TimeUsernameProblemLanguageResultExecution timeMemory
1369692gelastropodCollecting Stamps 5 (JOI26_stamps)C++20
100 / 100
1847 ms349708 KiB
#pragma GCC optimize("O3,inline")
#include <bits/stdc++.h>
using namespace std;
#define int long long

int D;
queue<int> ppl;
vector<int> T, ss, cp, cd, cnt1, cnt2, anss, md;
vector<vector<int>> adjlist, val1, val2, dd, pb, ff, ddi;

int dfs1(int n, int p) {
	ss[n] = 1;
	for (int i : adjlist[n]) {
		if (i == p || cd[i] != -1) continue;
		ss[n] += dfs1(i, n);
	}
	return ss[n];
}

int dfs2(int n, int p, int s) {
	for (int i : adjlist[n]) if (i != p && cd[i] == -1 && ss[i] > s) return dfs2(i, n, s);
	return n;
}

void dfs3(int n, int p) {
	int cc = dfs2(n, -1, dfs1(n, -1) / 2);
	cd[cc] = (p == -1 ? 0 : cd[p] + 1);
	cp[cc] = p;
	for (int i : adjlist[cc]) if (cd[i] == -1) dfs3(i, cc);
}

void dfs4(int n, int p, int src) {
	md[src] = max(md[src], dd[cd[src]][n]);
	ppl.push(n);
	val1[cd[src]][n] = min(val1[cd[src]][n], T[n] - dd[cd[src]][n]);
	val2[cd[src]][n] = min(val2[cd[src]][n], T[n] + dd[cd[src]][n]);
	for (int i = 0; i < adjlist[n].size(); i++) {
		if (adjlist[n][i] == p || cd[adjlist[n][i]] <= cd[src]) continue;
		pb[cd[src]][adjlist[n][i]] = (pb[cd[src]][n] == -1 ? i : pb[cd[src]][n]);
		val1[cd[src]][adjlist[n][i]] = val1[cd[src]][n], val2[cd[src]][adjlist[n][i]] = val2[cd[src]][n];
		dd[cd[src]][adjlist[n][i]] = dd[cd[src]][n] + 1;
		dfs4(adjlist[n][i], n, src);
	}
}

void proc(int n) {
	while (ppl.size()) ppl.pop();
	dfs4(n, -1, n);
	fill(cnt1.begin(), cnt1.begin() + adjlist[n].size() + 1 , 0);
	fill(cnt2.begin(), cnt2.begin() + adjlist[n].size() + 1, 0);
	for (int i = 0; i <= md[n]; i++) ff[i].clear();
	for (int i = 0; i <= md[n]; i++) ddi[i].clear();
	int tc1 = 0, tc2 = 0;
	while (ppl.size()) {
		int i = ppl.front(); ppl.pop();
		if (dd[cd[n]][i] > D) continue;
		cnt1[pb[cd[n]][i] + 1]++, tc1++;
		ddi[dd[cd[n]][i]].push_back(i);
		if (val1[cd[n]][i] < 0) cnt2[pb[cd[n]][i] + 1]++, tc2++;
		if (val1[cd[n]][i] >= 0 && val1[cd[n]][i] <= md[n]) ff[val1[cd[n]][i]].push_back(i);
	}
	for (int i = 0; i <= md[n]; i++) {
		for (int j : ff[i]) if (dd[cd[n]][j] <= D - i) cnt2[pb[cd[n]][j] + 1]++, tc2++;
		for (int j : ddi[i]) {
			if (val2[cd[n]][j] <= i) anss[j] += tc1 - (pb[cd[n]][j] == -1 ? 0 : cnt1[pb[cd[n]][j] + 1]);
			else anss[j] += tc2 - (pb[cd[n]][j] == -1 ? 0 : cnt2[pb[cd[n]][j] + 1]);
		}
		if (D - i < 0 || D - i > md[n]) continue;
		for (int j : ddi[D - i]) {
			cnt1[pb[cd[n]][j] + 1]--, tc1--;
			if (val1[cd[n]][j] <= i) cnt2[pb[cd[n]][j] + 1]--, tc2--;
		}
	}
}

signed main() {
	ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);
	int N, a, b; cin >> N >> D;
	ss.resize(N, 0);
	cp.resize(N, -1);
	cd.resize(N, -1);
	cnt1.resize(N, 0);
	cnt2.resize(N, 0);
	anss.resize(N, 0);
	md.resize(N, 0);
	anss.resize(N, 0);
	adjlist.resize(N, vector<int>());
	val1.resize(20, vector<int>(N, 1e15));
	val2.resize(20, vector<int>(N, 1e15));
	dd.resize(20, vector<int>(N, 0));
	pb.resize(20, vector<int>(N, -1));
	ff.resize(N, vector<int>());
	ddi.resize(N, vector<int>());
	for (int i = 0; i < N; i++) {
		cin >> a;
		T.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);
	}
	dfs3(0, -1);
	//for (int i : cp) cout << i << ' '; cout << '\n';
	//for (int i : cd) cout << i << ' '; cout << '\n';
	for (int i = 0; i < N; i++) proc(i);
	for (int i : anss) 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...