Untitled

 avatar
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