class Solution {
public:
void dfsCreateDirAdjList(int node, vector<bool>& marked, const vector<vector<int>>& adjList, vector<vector<int>>& dirAdjList)
{
marked[node] = true;
for(const int& adj : adjList[node])
{
if (!marked[adj])
{
dirAdjList[node].push_back(adj);
dfsCreateDirAdjList(adj, marked, adjList, dirAdjList);
}
}
}
void dfsPathToBob(int node, const vector<vector<int>>& dirAdjList, stack<int>& toBob, const int& bob, bool& foundBob)
{
toBob.push(node);
if (node == bob)
{
foundBob = true;
return;
}
for(const int& adj: dirAdjList[node])
{
dfsPathToBob(adj, dirAdjList, toBob, bob, foundBob);
if (foundBob)
{
return;
}
}
toBob.pop();
}
void dfsCheapestPath(int node, const vector<vector<int>>& dirAdjList, int& cost, const vector<int>& bobToZero, int level, const vector<int>& amount, unordered_set<int>& bobSeen, int& maxCost)
{
int curCost = 0;
level = (level < bobToZero.size()) ? level : bobToZero.size()-1;
bobSeen.insert(bobToZero[level]);
if (bobToZero[level] == node)
{
curCost = (amount[node]/2);
}
if (!bobSeen.contains(node))
{
curCost = amount[node];
}
cost += curCost;
if (dirAdjList[node].size() == 0)
{
maxCost = std::max(cost, maxCost);
}
else
{
for(int adj : dirAdjList[node])
{
dfsCheapestPath(adj, dirAdjList, cost, bobToZero, level+1, amount, bobSeen, maxCost);
}
}
if (level < bobToZero.size())
{
bobSeen.erase(bobToZero[level]);
}
cost -= curCost;
}
int mostProfitablePath(vector<vector<int>>& edges, int bob, vector<int>& amount) {
int n = edges.size() + 1;
// Create adjacency list for nodes in tree
vector<vector<int>> adjList = vector<vector<int>>(n, vector<int>{});
for(int ii=0; ii<edges.size(); ++ii)
{
int u = edges[ii][0];
int v = edges[ii][1];
adjList[u].push_back(v);
adjList[v].push_back(u);
}
// Create adjacency list with edges from root to leaves only.
vector<bool> marked(n, false);
vector<vector<int>> dirAdjList(n, vector<int>{});
dfsCreateDirAdjList(0, marked, adjList, dirAdjList);
// Create a vector of nodes from bob's leaf to zero.
vector<int> bobToZero{};
stack<int> toBob{};
bool foundBob = false;
dfsPathToBob(0, dirAdjList, toBob, bob, foundBob);
bobToZero.reserve(toBob.size());
while(!toBob.empty())
{
bobToZero.push_back(toBob.top());
toBob.pop();
}
// Alice walks all paths from root to leaves, while simulating Bob
// walking from Bob's leaf to root.
unordered_set<int> bobSeen;
int maxCost = std::numeric_limits<int>::min();
int cost = 0;
dfsCheapestPath(0, dirAdjList, cost, bobToZero, 0, amount, bobSeen, maxCost);
return maxCost;
}
};