Sum of all paths between all pairs of nodes in a tree

Question

I'm trying to solve a competitive programming problem. It basically gives a undirected graph (tree-like: no multiple paths between two nodes...) and asks for the sum of all possible paths between any pair of nodes in the graph (each path must be counted only once, in other words, if you have already counted the path from A to B, you shouldn't count the path from B to A).

To solve this, I've tried to remove each edge (separating the tree in 2 components) and count the number of vertices in one of the components (the count of nodes of the second component can be derived from the count found in the first).

The total weight that a edge will contribute is: (the number of paths that pass by it = first_component_size * second_component_size) * (the weight of the edge).

Input description:

The input is composed of several instances (the number of instances is given on the first line of input).

The first input row of each instance contains an integer N (1 ≤ N ≤ 10000), representing the number of nodes. The nodes are enumerated from 1 to N.

Each one of the following N-1 rows contains three integers A, B and C (1 ≤ A, B ≤ N, 1 ≤ C ≤ 50), indicating that the edge that connects the nodes A and B has length C.

Output description:

For each instance print the sum of lenghts of the paths that connects all the pairs of nodes. The answer must be in MOD 1300031.

#include <cstdio>
#include <vector>
#include <utility>

const int MOD = 1300031;

using edge = std::pair<int, int>;
using adj_list   = std::vector<edge>;
using graph  = std::vector<adj_list> ;


struct myEdge {
    int from, to, weight;
};

std::vector<bool> visited;
int count_nodes_on_component(graph &g, int v) {
    int count_of_nodes = 1; //1 for itself
    adj_list &adj = g[v];
    for(auto &e : adj) {
        if(!visited[e.first]) {
            visited[e.first] = true;
            count_of_nodes += count_nodes_on_component(g, e.first);
            visited[e.first] = false;
        }
    }
    return count_of_nodes;
}


int main ( void ) {
    int test_cases;
    scanf("%d", &test_cases);
    while (test_cases--) {
        int total_nodes_in_graph;
        scanf("%d", &total_nodes_in_graph);
        
        visited.resize(total_nodes_in_graph + 1);
        
        graph g(total_nodes_in_graph + 1);
        std::vector<myEdge> vector_all_edges;
        for (int i = 1, node1, node2, weight; i < total_nodes_in_graph; ++i) {
            scanf("%d %d %d", &node1, &node2, &weight);
            g[node1].push_back( {node2, weight} );
            g[node2].push_back( {node1, weight} );
            vector_all_edges.push_back( {node1, node2, weight} );
        }
        
        long long total_sum = 0;
        for (auto &e : vector_all_edges) {
            visited[e.from] = visited[e.to] = true;
            int cc = count_nodes_on_component(g, e.to);
            total_sum = (total_sum + cc * (total_nodes_in_graph - cc) * e.weight) % MOD;
            visited[e.from] = visited[e.to] = false;
        }
        
        printf("%lld\n", total_sum);
    }
}

This code takes a lot of time with huge graphs. Is it possible to exploit something to make it run faster?

Answer 1

\$O(n^2)\$ time complexity

Your current algorithm has \$O(n^2)\$ time complexity because for each edge, you count the number of vertices on one side of the edge, which is a linear time operation.

However, the program is pretty fast. I ran it on three large test cases. The first was 10000 nodes arranged in a line, the second was 10000 nodes arranged in a star, and the third was 10000 nodes arranged in a star, but with the edges listed with reverse direction. Here were the results:

Line : 1.15 sec
Star : 0.01 sec
Star2: 2.30 sec

The first star was quick because the program always counted the outside (i.e. 1 node always). The second star was slow because the program always counted the inside (n-1 nodes always). So you could consider these the best and worst possible test cases.

Alternative strategy

The problem can be solved in \$O(n)\$ time. Instead of picking a random edge and having to count one side, if you picked an edge connected to a leaf node, you automatically know that the count on the leaf side is 1. There is always at least 1 leaf node in the graph, because the graph is a spanning tree with no cycles. If you then remove the leaf node from the graph, there must still be at least 1 leaf node, so you can keep finding a leaf node and removing it until there are no nodes left.

A small complication arises after you find the first leaf node and remove it from the tree. The node it was attached to must now remember that it was connected to the leaf, in order to keep the distance calculations correct. In other words, if leaf A is attached to node B and leaf A is removed, then node B must now count as 2 nodes. If node B is later removed, it must add its 2 nodes to some other node, etc.

So the pseudocode is this:

1. Find a node with 1 edge only. Call this node A and the node it is attached to B.
2. Add to the total distance this amount: A.count * (totalNodes - A.count) * edge.weight.
3. Add A.count to B.count.
4. Remove the edge from A and B.
5. Go back to step 1. If B only has 1 edge, use B as the next node.

You can find all the leaf nodes in \$O(n)\$ time because:

1. In the beginning there are a fixed number of leaf nodes.
2. If you scan the nodes array from beginning to end, you will find these initial leaf nodes in \$O(n)\$ time.
3. Any other nodes that become leaf nodes are found in step 5 above and can be handled without scanning.

Similarly, finding and removing edges can be done in amortized \$O(n)\$ time if you handle the edge lists correctly. One key point is that even though you need to search through edge lists, there are only \$n\$ total edges, so that total time spent searching through edge lists is limited to \$O(n)\$.

Sample \$O(n)\$ program

Here is the program I wrote in C with the alternative strategy. It solved all three test cases in 0.01 seconds.

Note: I did edge removal in a somewhat nonstandard way, where I marked the edge invalid instead of removing it from the edge list. I wrote another version where I kept the edges in a doubly linked circular list and I actually removed the edge from the list when needed. Both versions are just as fast, so if anyone wants to see the other version, just ask.

#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>

const int MOD = 1300031;

typedef struct Edge {
    struct Edge *next;
    int          to;
    int          weight;
} Edge;

typedef struct Node {
    int         numConnected;
    int         numEdges;
    Edge       *edges;
} Node;

int main(void)
{
    int   numTests = 0;
    Node *nodes    = NULL;
    Edge *edges    = NULL;

    scanf("%d", &numTests);
    for (int i=0; i<numTests; i++) {
        int numNodes = 0;
        int e        = 0;

        scanf("%d", &numNodes);
        nodes = (Node *) calloc(numNodes, sizeof(*nodes));
        edges = (Edge *) calloc(2*numNodes, sizeof(*edges));

        // Read in list of edges.  For each edge, we create two Edge
        // structures and add one to each endpoint Node.
        for (int j = numNodes - 1; j > 0; j--) {
            int from   = 0;
            int to     = 0;
            int weight = 0;
            scanf("%d %d %d", &from, &to, &weight);
            from--;
            to--;

            edges[e].to       = to;
            edges[e].next     = nodes[from].edges;
            edges[e].weight   = weight;
            nodes[from].edges = &edges[e];
            nodes[from].numEdges++;
            e++;

            edges[e].to       = from;
            edges[e].next     = nodes[to].edges;
            edges[e].weight   = weight;
            nodes[to].edges   = &edges[e];
            nodes[to].numEdges++;
            e++;
        }

        int      search = 0;
        int      cur    = -1;
        uint64_t dist   = 0;

        for (int nodesLeft = numNodes-1; nodesLeft > 0; nodesLeft--) {
            Edge *edge;

            // Either use the node known to have 1 edge, or find the next
            // node that has 1 edge.
            if (cur == -1) {
                while (nodes[search].numEdges != 1) {
                    search++;
                    if (search >= numNodes) {
                        printf("ERROR, reached end of nodes\n");
                        exit(1);
                    }
                }
                cur = search++;
            }

            // Find the one valid edge in the list of edges.
            for (edge = nodes[cur].edges; edge != NULL; edge = edge->next) {
                if (edge->to != -1)
                    break;
            }
            if (edge == NULL) {
                printf("ERROR, reached end of edges\n");
                exit(1);
            }

            // The edge is going from cur <-> to, where cur is a leaf node.
            // We can calculate the distance used by the edge because we know
            // how many nodes are attached to cur.  That number is 1 for cur
            // itself plus nodes[cur].numConnected which is the number of
            // previous leaf nodes that we removed that are connected to cur.
            int to       = edge->to;
            int curNodes = nodes[cur].numConnected + 1;

            dist += (numNodes - curNodes) * curNodes * edge->weight;

            // Since we are removing the leaf node cur, we add its nodes
            // to the node it is attached to.
            nodes[to].numConnected += curNodes;

            // If the new node has 1 edge, use it, otherwise set cur to
            // negative to search for another node with 1 edge.
            if (--nodes[to].numEdges == 1)
                cur = to;
            else
                cur = -1;
            nodes[cur].numEdges = 0;

            // Delete other side of edge by marking its to field as -1.
            // This is needed so that when we search through the other
            // node's edges for its one remaining edge, we will skip this one.
            // We can find the paired edge by looking at its array index
            // because always created edges in consecutive pairs.
            int edgePairIndex = (edge - &edges[0]) ^ 0x1;
            edges[edgePairIndex].to = -1;
        }

        // printf("%lld\n", dist);
        printf("%lld\n", dist % MOD);
        free(edges);
        free(nodes);
    }
}

Answer 2

To add, you can avoid using recursion (and avoid running out of stack memory) using dfs (with an actual stack) and marking each position by their post-order traversal. Then perform the operations on the nodes by order of their post-order traversal.

Answer 3

For the given problem, my answer would be similar to Faizan's one. But if you need to calculate the sum of distance of a node to all the other nodes you can follow this procedure.

I have performed dfs twice and stored the sum of distances of the node. Finally I just traverse the array and add half of the stored value to the final answer. Time complexity is again O(n).

#include<bits/stdc++.h>
using namespace std;
#define mp make_pair
typedef long long ll;
#define MOD 1300031
vector<vector<pair<ll,ll>>> g;
vector<ll> sum1,sum2,cnt;
vector<int> visited; 
ll n;
void dfs1(ll node)
{
    visited[node] = 1;
    for(int i = 0;i<g[node].size();i++)
    {
        ll x = g[node][i].first, w = g[node][i].second;
        if(visited[x]) continue;
        dfs1(x);
        sum1[node] += sum1[x] + cnt[x]*w;
        cnt[node] += cnt[x];
    }
    cnt[node]++;
}
void dfs2(ll node)
{
    if(node == 1)
    {
        sum2[node] = sum1[node];
    }
    visited[node] = 1;
    for(int i=0;i<g[node].size();i++)
    {
        ll x = g[node][i].first, w = g[node][i].second;
        if(visited[x]) continue;
        sum2[x] = sum1[x] + sum2[node] - sum1[x] - cnt[x]*w + w*(n-cnt[x]);
        dfs2(x);
    }
}
int main()
{
    int t,x,y,w;
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cin>>t;
    while(t--)
    {
        cin>>n;
        g.resize(n+1);sum1.resize(n+1);sum2.resize(n+1);
        visited.resize(n+1);cnt.resize(n+1);
        for(int i=0;i<n-1;i++)
        {
            cin>>x>>y>>w;
            g[x].push_back(mp(y,w));
            g[y].push_back(mp(x,w));
        }
        dfs1(1);
        visited.clear();
        visited.resize(n+1);
        dfs2(1);
        double ans = 0;
        for(int i=1;i<=n;i++)
        {
            ans+=(double)sum2[i]/2;
        }
        ll fans = (ll)ans;
        cout<<fans%MOD<<"\n";
        sum1.clear();sum2.clear();visited.clear();cnt.clear();g.clear();
    }

}

Answer 4

The method you used here is O(n^2) in complexity.. Here is my AC solution which is O(n) in comlexity and needs just on dfs to find total sum of all distinct pair of path.

/*Idea used in DFS---
while returning from dfs function for each node subtract subsize of that  
node(say 's') from total number of nodes in tree let's say 'u'.
Now, result+=(s*(u-s)*cost);(cost=cost of edge,connecting that vertex to it's parent*/

    //lovetocode.cpp
    #include <bits/stdc++.h>
    using namespace std;

    #define ll long long
    #define pb push_back
    #define fr_z(start,end) for(int z=start;z<end;z++)
    #define w while
    #define mod 1300031
    #define mp make_pair
    #define fa_io std::ios::sync_with_stdio(false) //just for fast i/o

    #define N 10100 //size of graph
    #define root 1 //root of dfs tree

    vector< pair<int,int> > g[N]; //pair.first=node,pair.second-cost
    int subsize[N],n; //subsize[N] contains subsize of each node in dfs tree
    ll res;

    void dfs(int node,int cost,int parent)
    {
        subsize[node]=1; //subsize of each node starts with one
        //cout<<node<<' '<<parent<<'\n';
        for(auto it:g[node])
        {
            int next=it.first;
            int c=it.second;
            if(next!=parent)
            {
                dfs(next,c,node);
                subsize[node]+=subsize[next]; //subsize[node]+=(subsize[of all of it's child node)
            }
        }
        res=( (((((ll)(n-subsize[node])%mod) * ((ll)subsize[node]%mod))%mod) * (((ll)cost)%mod))%mod + res%mod)%mod; //this is the same formula as given by you above, so THANX for it :)
    }

    int main()
    {
        fa_io;
        cin.tie(NULL);
        int t,a,b,c;
        cin>>t; //no. of testcases
        while(t--)
        {
            cin>>n; //no of nodes
            fr_z(1,n)
            {
                cin>>a>>b>>c;
                g[a].pb(mp(b,c));
                g[b].pb(mp(a,c)); //hope you get this in case if you don't comment or text me i'll explain
            }
            res=0;
            dfs(root,0,-1);
            fr_z(1,n+1)
                g[z].clear();
            cout<<res<<'\n';
        }

        return 0;
    }