#include <iostream>
#include <vector>
using namespace std;
typedef long long int ll;
ll dfs(int S, int parent, vector<vector<int>> & adjList, vector<ll> & DP) {
for (int v: adjList[S]) {
if (v == parent)
continue;
DP[S] += dfs(v, S, adjList, DP);
}
// DP[S] contains the total population in the subtree of S
return DP[S];
}
int LocateCentre(int N, int P[], int S[], int D[]) {
vector<ll>DP(N);
vector<vector<int>> adjList(N, vector<int>());
for (int i = 0; i < N-1; i++) {
// we convert the graph format to an adjacency list
adjList[S[i]].emplace_back(D[i]);
adjList[D[i]].emplace_back(S[i]);
}
for (int i = 0; i < N; i++)
DP[i] = P[i]; // we initialise the DP so that it contains each city's population
dfs(0, -1, adjList, DP);
// we choose 0 originally as our city with lowest maximum congestion and find its largest congestion
int chosenCity = 0;
ll chosenCityMaxCongestion = 0;
for (int rootNeighbour: adjList[0]) {
chosenCityMaxCongestion = max(chosenCityMaxCongestion, DP[rootNeighbour]);
}
// we iterate over all possible chosen cities and compare them with the current record
for (int currentCity = 1; currentCity < N; currentCity++) {
// before comparing with the congestion coming from children nodes, we will take as the maximum the congestion travelling from the currentCity's parent to the currentCity
ll currentCityMaxCongestion = DP[0] - DP[currentCity];
// we now iterate over the currentCity's children nodes to see if the congestion from a child to the currentCity is higher than that from the parent
for (int v: adjList[currentCity]) {
if (DP[v] < DP[currentCity])
currentCityMaxCongestion = max(currentCityMaxCongestion, DP[v]);
}
// if the city's maximum congestion is lower than the current record, we update it
if (currentCityMaxCongestion < chosenCityMaxCongestion) {
chosenCityMaxCongestion = currentCityMaxCongestion;
chosenCity = currentCity;
}
}
return chosenCity;
}
