Submission #1871747

---

Source Code Expand

#include "iostream"
#include "climits"
#include "list"
#include "queue"
#include "stack"
#include "set"
#include "functional"
#include "algorithm"
#include "string"
#include "map"
#include "unordered_map"
#include "iomanip"
#include "cmath"

using namespace std;

const long long int MOD = 1000000007;

long long int N, M, K, H, W, L, R;

class Lowest_Common_Ancestor {
	vector<int>depth;
	vector<int>parent[25];
	list<int>edge[1000001];
	int node;
public:
	Lowest_Common_Ancestor(int num) {
		node = num;
		num++;
		depth.resize(num);
		for (int i = 0; i < 25; i++)parent[i].resize(num);
	}
	void Add_Edge(int a, int b) {
		edge[a].push_back(b);
		edge[b].push_back(a);
		return;
	}
	void Update() {
		queue<int>QQ;
		depth[1] = 0;
		for (int i = 2; i <= node; i++) depth[i] = INT_MAX;
		QQ.push(1);
		while (!QQ.empty()) {
			int c = QQ.front();
			for (auto i : edge[c]) {
				if (depth[i] > depth[c] + 1) {
					depth[i] = depth[c] + 1;
					QQ.push(i);
				}
			}
			QQ.pop();
		}
		parent[0][1] = -1;
		for (int i = 2; i <= node; i++) {
			for (auto j : edge[i]) {
				if (depth[i] - 1 == depth[j]) {
					parent[0][i] = j;
					break;
				}
			}
		}
		for (int i = 0; i < 24; i++) {
			for (int j = 1; j <= node; j++) {
				if (parent[i][j] < 0)parent[i + 1][j] = -1;
				else parent[i + 1][j] = parent[i][parent[i][j]];
			}
		}
		return;
	}
	int LCA(int u, int v) {
		if (depth[u] > depth[v])swap(u, v);
		for (int i = 0; i < 25; i++) {
			if ((depth[v] - depth[u]) >> i & 1) {
				v = parent[i][v];
			}
		}
		if (u == v)return u;
		for (int i = 24; i >= 0; i--) {
			if (parent[i][v] != parent[i][u]) {
				u = parent[i][u];
				v = parent[i][v];
			}
		}
		return parent[0][u];
	}
	int Dist(int u, int v) {
		return depth[u] + depth[v] - depth[LCA(u, v)] * 2;
	}
};

int main() {
	ios::sync_with_stdio(false);
	cin.tie(0);

	cin >> N >> K;
	Lowest_Common_Ancestor lca(N + 1);
	vector<vector<int>>v(N + 1);
	vector<vector<int>>depth(N);
	vector<pair<int, int>>place(N + 1);
	vector<vector<int>>edge(N + 1);
	vector<int>flag(N + 1);
	vector<int>check;
	for (int i = 1; i <= N; i++) {
		cin >> M;
		if (M&&i!=K) {
			check.push_back(i);
		}
	}
	M = 0;
	for (int i = 1; i < N; i++) {
		cin >> L >> R;
		edge[L].push_back(R);
		edge[R].push_back(L);
		lca.Add_Edge(L, R);
		lca.Add_Edge(R, L);
	}
	queue<int>Q;
	Q.push(K);
	for (int i = 1; i <= N; i++) {
		sort(edge[i].begin(), edge[i].end());
	}
	while (!Q.empty()) {
		int cn = Q.front();
		Q.pop();
		int box = 1;
		for (auto i : edge[cn]) {
			if (v[i].empty()&&i!=K) {
				Q.push(i);
				v[i] = v[cn];
				v[i].push_back(box);
				box++;
			}
		}
	}
	vector<vector<int>>search;
	for (auto i : check) {
		v[i].push_back(0);
		v[i].push_back(i);
		search.push_back(v[i]);
	}
	search.push_back({ 0,(int)K });
	sort(search.begin(), search.end());
	search.push_back({ 0,(int)K });
	lca.Update();
	for (int i = 1; i < search.size(); i++) {
		M += lca.Dist(search[i].back(), search[i - 1].back());
	}
	cout << M << endl;
	return 0;
}

Submission Info

Judge Result