#include <bits/stdc++.h>
#include "traffic.h"
using namespace std;
typedef long long int ll;
vector<vector<int>> graph;
vector<ll> cost;
vector<bool> visited;
vector<ll> DP;
int arenaCity;
ll traffic (int city) {
if (DP[city] != -1) return DP[city];
DP[city] = cost[city];
visited[city] = true;
for (int road: graph[city]) {
if (!visited[road]) DP[city] += traffic(road);
}
return DP[city];
}
void untraffic (int city, int cityMaxDiff) {
int bestCity = city;
ll bestDiff = cityMaxDiff, maxCongestion = DP[city];
// cerr << city << ' ' << cityMaxDiff << '\n';
for (int road: graph[city]) {
// if (DP[road] < maxCongestion) continue;
ll diff = maxCongestion - DP[road];
// cerr << diff << '\n';
if (diff < bestDiff) {
bestCity = road;
bestDiff = diff;
}
}
if (bestCity != city) {
DP[city] -= DP[bestCity];
DP[bestCity] = maxCongestion;
arenaCity = bestCity;
ll congestion = 0;
for (auto road: graph[bestCity]) {
congestion = max(congestion, DP[road]);
}
untraffic(bestCity, congestion);
}
}
int LocateCentre(int N, int pp[], int S[], int D[]) {
graph = vector<vector<int>> (N);
cost.clear();
DP.assign(N, -1);
for (int i = 0; i < N-1; i++) {
graph[S[i]].push_back(D[i]);
graph[D[i]].push_back(S[i]);
}
for (int i = 0; i < N; i++) {
cost.push_back(pp[i]);
}
visited.assign(N, false);
ll congestion = traffic(0);
congestion = 0;
for (auto road: graph[0]) {
// cerr << road << '\n';
congestion = max(congestion, DP[road]);
}
// cerr << congestion << '\n';
arenaCity = 0;
untraffic(0, congestion);
return arenaCity;
}
// int main() {
// int N = 5;
// int pp[5] = {10, 10, 10, 10, 10};
// int S[4] = {0, 1, 2, 3};
// int D[4] = {1, 2, 3, 4};
// cout << LocateCentre(N, pp, S, D) << '\n';
// for (int i: DP) cout << i << ' ';
// cout << '\n';
// }
