#include "factories.h"
#include <bits/stdc++.h>
using namespace std;
const int MAXN = 5e5 + 5;
vector<pair<int, int> > g[MAXN];
int Subtree[MAXN], lastCnt[MAXN], centroidParent[MAXN];
long long distFromRoot[MAXN], ans[MAXN];
bool isTaken[MAXN];
int DP[20][MAXN], depth[MAXN];
int cnt = 0;
void FindDist(int node, int par) {
DP[0][node] = par;
for(int i = 0; i < g[node].size(); i++) {
if(g[node][i].first != par) {
distFromRoot[g[node][i].first] = distFromRoot[node] + g[node][i].second;
depth[g[node][i].first] = depth[node] + 1;
FindDist(g[node][i].first, node);
}
}
}
int SubtreeDfs(int node, int par) {
Subtree[node] = 1;
for(int i = 0; i < g[node].size(); i++) {
if(!isTaken[g[node][i].first] && !isTaken[node] && g[node][i].first != par) {
Subtree[node] += SubtreeDfs(g[node][i].first, node);
}
}
return Subtree[node];
}
int FindCentroid(int node, int par, int totalSize) {
for(int i = 0; i < g[node].size(); i++) {
if(!isTaken[g[node][i].first] && !isTaken[node] && g[node][i].first != par) {
if(Subtree[g[node][i].first] > totalSize / 2) {
return FindCentroid(g[node][i].first, node, totalSize);
}
}
}
return node;
}
void BuildCentroid(int centroid, int par) {
int siz = SubtreeDfs(centroid, -1);
centroid = FindCentroid(centroid, par, siz);
isTaken[centroid] = true;
if(par == -1)
centroidParent[centroid] = centroid;
else
centroidParent[centroid] = par;
for(int i = 0; i < g[centroid].size(); i++) {
if(!isTaken[g[centroid][i].first]) {
BuildCentroid(g[centroid][i].first, centroid);
}
}
}
void Init(int N, int A[], int B[], int D[]) {
for(int i = 0; i < MAXN; i++) {
isTaken[i] = false;
Subtree[i] = 0;
distFromRoot[i] = 0;
centroidParent[i] = i;
lastCnt[i] = 0;
ans[i] = 1e15;
}
cnt = 0;
for(int i = 0; i < N - 1; i++) {
g[A[i]].push_back({B[i], D[i]});
g[B[i]].push_back({A[i], D[i]});
}
distFromRoot[0] = 0;
depth[0] = 0;
FindDist(0, -1);
BuildCentroid(0, -1);
for(int i = 1; i < 20; i++) {
for(int j = 0; j < N; j++) {
if(DP[i - 1][j] == -1) {
DP[i][j] = -1;
}
else {
DP[i][j] = DP[i - 1][DP[i - 1][j]];
}
}
}
/*cout << "Centroid Parents\n";
for(int i = 0; i < N; i++) {
cout << centroidParent[i] << " ";
}
cout << "\n";
cout << "distFromRoot\n";
for(int i = 0; i < N; i++) {
cout << distFromRoot[i] << " ";
}
cout << "\n";
cout << "Normal parents\n";
for(int i = 0; i < N; i++) {
cout << DP[0][i] << " ";
}
cout << "\n";
cout << "Binary lifting array\n";
for(int i = 0; i < 20; i++) {
for(int j = 0; j < N; j++) {
cout << DP[i][j] << " ";
}
cout << "\n";
}*/
}
int LCA(int x, int y) {
if(depth[x] < depth[y])
swap(x, y);
int diff = depth[x] - depth[y];
/*cout << "IN LCA\n";
cout << "x = " << x << ", y = " << y << ", diff = " << diff << "\n";
cout << "(1 << 2) = " << (1 << 2) << "\n";*/
for(int i = 19; i >= 0; i--) {
if(diff >= (1 << i)) {
diff -= (1 << i);
x = DP[i][x];
}
}
//cout << "x = " << x << ", y = " << y << "\n";
if(x == y) {
//cout << "Answer = " << x << "\n";
return x;
}
for(int i = 19; i >= 0; i--) {
if(DP[i][x] != DP[i][y]) {
x = DP[i][x];
y = DP[i][y];
}
}
//cout << "Answer = " << DP[0][x] << "\n";
return DP[0][x];
}
void Update(int node) {
ans[node] = 0;
lastCnt[node] = cnt;
int initial = node;
while(centroidParent[node] != node) {
node = centroidParent[node];
//cout << "Update function, updating " << initial << "\n";
//cout << "initial = " << initial << ", node = " << node << ", distance = " << distFromRoot[node] + distFromRoot[initial] - 2 * distFromRoot[LCA(initial, node)] << "\n";
if(lastCnt[node] != cnt) {
ans[node] = distFromRoot[node] + distFromRoot[initial] - 2 * distFromRoot[LCA(initial, node)];
lastCnt[node] = cnt;
}
else {
ans[node] = min(ans[node], distFromRoot[node] + distFromRoot[initial] - 2 * distFromRoot[LCA(initial, node)]);
}
//cout << "ans[node] = " << ans[node] << "\n";
}
}
long long Query2(int node) {
if(lastCnt[node] == cnt)
return 0;
int initial = node;
long long res = 1e15;
while(centroidParent[node] != node) {
node = centroidParent[node];
//cout << "Query function, querying " << initial << "\n";
//cout << "LCA(initial, node) = " << LCA(initial, node) << "\n";
//cout << "initial = " << initial << ", node = " << node << ", distance = " << distFromRoot[node] + distFromRoot[initial] - 2 * distFromRoot[LCA(initial, node)] << "\n";
if(lastCnt[node] == cnt) {
res = min(res, distFromRoot[node] + distFromRoot[initial] - 2 * distFromRoot[LCA(initial, node)] + ans[node]);
}
//cout << "res = " << res << "\n";
}
return res;
}
long long Query(int S, int X[], int T, int Y[]) {
cnt++;
for(int i = 0; i < S; i++) {
Update(X[i]);
}
/*cout << "Outputting ans array\n";
for(int i = 0; i < N; i++) {
cout << ans[i] << " ";
}
cout << "\n";
cout << "Outputting lastCnt array\n";
for(int i = 0; i < N; i++) {
cout << lastCnt[i] << " ";
}
cout << "\n";*/
long long res = 1e15;
for(int i = 0; i < T; i++) {
res = min(res, Query2(Y[i]));
}
return res;
}
Compilation message
factories.cpp: In function 'void FindDist(int, int)':
factories.cpp:15:22: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
15 | for(int i = 0; i < g[node].size(); i++) {
| ~~^~~~~~~~~~~~~~~~
factories.cpp: In function 'int SubtreeDfs(int, int)':
factories.cpp:26:22: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
26 | for(int i = 0; i < g[node].size(); i++) {
| ~~^~~~~~~~~~~~~~~~
factories.cpp: In function 'int FindCentroid(int, int, int)':
factories.cpp:35:22: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
35 | for(int i = 0; i < g[node].size(); i++) {
| ~~^~~~~~~~~~~~~~~~
factories.cpp: In function 'void BuildCentroid(int, int)':
factories.cpp:53:22: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
53 | for(int i = 0; i < g[centroid].size(); i++) {
| ~~^~~~~~~~~~~~~~~~~~~~
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Incorrect |
35 ms |
26988 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Incorrect |
20 ms |
26604 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Incorrect |
35 ms |
26988 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
