Untitled
unknown
plain_text
a year ago
1.6 kB
22
Indexable
#include <vector>
using namespace std;
vector<int> countReverseEdges(int g_nodes, vector<int> g_from, vector<int> g_to) {
int n = g_nodes;
vector<vector<pair<int,int>>> adj(n+1);
for (int i = 0; i < g_from.size(); ++i) {
int u = g_from[i], v = g_to[i];
adj[u].push_back({v, 0}); // original edge, cost 0
adj[v].push_back({u, 1}); // reverse edge, cost 1 (needs reversal if going from v->u)
}
vector<int> res(n+1);
vector<int> visited(n+1, 0);
// First DFS: count reversals needed to reach all from node 1
function<void(int, int)> dfs1 = [&](int u, int parent) {
for (auto &p : adj[u]) {
int v = p.first, cost = p.second;
if (v != parent) {
res[1] += cost;
dfs1(v, u);
}
}
};
dfs1(1, 0);
// Second DFS: reroot and propagate answer to all nodes
function<void(int, int)> dfs2 = [&](int u, int parent) {
for (auto &p : adj[u]) {
int v = p.first, cost = p.second;
if (v != parent) {
if (cost == 0) // edge u->v, when rerooting, needs to be reversed
res[v] = res[u] + 1;
else // edge v->u, when rerooting, does NOT need to be reversed
res[v] = res[u] - 1;
dfs2(v, u);
}
}
};
dfs2(1, 0);
// The result for 1-based nodes
vector<int> finalRes(n);
for (int i = 1; i <= n; ++i) finalRes[i-1] = res[i];
return finalRes;
}
Editor is loading...
Leave a Comment