#include "traffic.h"
#include <bits/stdc++.h>
using namespace std;
struct Node
{
vector<Node*> neigh;
int size_up;
int population;
Node* up;
};
int sizeOfSubtree(Node *t, Node *p)
{
int c = t->population;
t->up = p;
for(auto i : t->neigh)
{
if (i==p) continue;
c += sizeOfSubtree(i,t);
}
if(t->up!=nullptr)
t->size_up = c;
else
t->size_up = 0;
return c;
}
int maxSize(Node *t,int max_size)
{
int c=t->population;
int m = INT_MIN;
for(auto i : t->neigh)
{
if (i==t->up) continue;
m = max(m,i->size_up);
c += i->size_up;
}
return max(m,max_size-c);
}
int LocateCentre(int N, int pp[], int S[], int D[]) {
vector<Node> nodes(N);
for(int i = 0; i < N; i++)
{
nodes[i].population = pp[i];
}
for(int i = 0; i < N-1; i++)
{
nodes[S[i]].neigh.push_back(&nodes[D[i]]);
nodes[D[i]].neigh.push_back(&nodes[S[i]]);
}
int total_size = sizeOfSubtree(&nodes[0],nullptr);
int m = INT_MAX;
int midx = -1;
for(int i = 0; i < N; i++)
{
int a = maxSize(&nodes[i],total_size);
//cout << a << " " << i << "\n";
if (m > a)
{
m = a;
midx = i;
}
}
return midx;
}
