Sum of maximum flow between every pair of vertices of a tree

Given a undirected tree with N vertices numbered from 1 to N. Each edge tree has some capacity. Find the sum of maximum flow between all possible pair of vertices. There exist only one way to go between any two vertices.
Find the sum of maximum flow between all possible pair of vertices.

For eg: In given tree with 3 edges
1 2 5
2 3 6
Edges between node 1 and node 2 with capacity 5, node 2 and node 3 with capacity 6.
Output - 32

(1,2) = (2,1) = 5
(1,3) = (3,1) = 5
(2,3) = (3,2) = 6
Therefore output is (5+5+6)*2 = 32

My Approach-

1. Sort edges based on edge_capacity

2. While edge_list is not empty: remove edge with minimum capacity

   • count number of nodes on left and right of this edge. Do DFS for node count
   • Add (left_count * right_count * edge_capacity) to the answer.

3. return answer*2.
   

Time complexity is O(n²). This solution gives TLE.
How can we further reduce the time complexity?

My code-

def dfs(node):
    count = 1
    visited = set()
    stack = [node]
    while stack:
        vertex = stack.pop()
        if vertex not in visited:
            visited.add(vertex)
            stack.extend(set(nodes[vertex]) - visited)
    return len(visited)

for _ in range(int(raw_input())):   # Iterate for multiple test cases
    MOD = 1000000007
    edges = []
    n = int(raw_input())            # number of vertices
    nodes = [set() for _ in range(n)]
    for __ in range(n-1):           # read input for number of edges
        edges.append(map(int, raw_input().split()))
        nodes[edges[-1][0]-1].add(edges[-1][1]-1)    
        nodes[edges[-1][1]-1].add(edges[-1][0]-1)
        edges[-1][0]-=1; edges[-1][1]-=1; 
    edges.sort(key=lambda x: x[2])

    answer = 0
    for i in range(len(edges)):
        weight = edges[i][2]
        x, y = edges[i][0], edges[i][1]
        nodes[x].remove(y)
        nodes[y].remove(x)
        left_count = dfs(x)
        right_count = dfs(y)
        answer += ((left_count*right_count*weight)%MOD)
    print (answer*2)%MOD

Link to the original problem- Spoj-Flow on Tree

Updates

Constraints-

1. Number of test cases - 10
2. 1 <= N <= 10⁵ (number of vertices in each test case)
3. Capacities of each edge will be non-negative and no more than 10⁶.
4. The total number of vertices among all test cases will be less than 5*10⁵.

Solution 1:

Instead of running two new DFSes each time to count subtree sizes, run a smarter DFS just once that computes, for each edge uv, the number of nodes in the subtree rooted at u that is formed when the edge uv is deleted. (Note that you will need to compute this value for both uv and vu.)

For a way to compute these node counts in linear time, see this nice answer.