#include <iostream>
#include <vector>
using namespace std;
typedef long long int ll;
ll dfs(int S, int parent, vector<vector<int>> &adjList,
vector<ll> &subtreePopulation) {
for (int neighbouringCity : adjList[S]) {
if (neighbouringCity == parent)
continue;
subtreePopulation[S] += dfs(neighbouringCity, S, adjList,
subtreePopulation);
}
// subtreePopulation[S] contains the total population in the subtree of S
return subtreePopulation[S];
}
int chooseCity(vector<vector<int>> &adjList, vector<ll> &subtreePopulation) {
int N = (int) adjList.size();
// 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,
subtreePopulation[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 = subtreePopulation[0]
- subtreePopulation[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 neighbouringCity : adjList[currentCity]) {
if (subtreePopulation[neighbouringCity]
< subtreePopulation[currentCity]) // if this is not the case, neighbouringCity is the parent of currentCity
currentCityMaxCongestion = max(currentCityMaxCongestion,
subtreePopulation[neighbouringCity]);
}
// if the city's maximum congestion is lower than the current record, we update it
if (currentCityMaxCongestion < chosenCityMaxCongestion) {
chosenCityMaxCongestion = currentCityMaxCongestion;
chosenCity = currentCity;
}
}
return chosenCity;
}
int LocateCentre(int N, int P[], int S[], int D[]) {
vector<ll> subtreePopulation(N);
vector<vector<int>> adjList(N);
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++)
subtreePopulation[i] = P[i]; // we initialise the subtreePopulation vector so that it contains each city's population
dfs(0, -1, adjList, subtreePopulation);
return chooseCity(adjList, subtreePopulation);
}
