Count number of pairs of nodes in undirected graph such that W - L >= K

Question

Question:

Given an undirected graph of N nodes and N-1 edges. The length of all edges are 1. Each edge i has a weight of Wi. (0 <= Wi <= 10.000)

The graph is guaranteed to not have any loops. So there's always one (and only) shortest path between 2 nodes.

Take a pair of node (u, v) from the graph:

• l is the length of shortest the path between the 2 nodes
• w is the edge with the largest weight in the shortest path between the 2 nodes

Given the number K, count the number of pair (u, v) from the graph such that w - l >= K

Example:

N = 3, K = 1
Edges:
1 - 2 - 3
1 - 3 - 2

(Edges are described as: u - v - w)

Answer: 6. All the possible pairs are: (1, 2), (1, 3), (2, 1), (2, 3), (3, 1), (3, 2)

Brute Force Solution (N < 100):

Go over all pairs of (u, v), find the shortest path along side with w and l. Increase the answer if w - l >= K.

Time complexity: O(N^3)

P/S: I have tried and succeeded the brute force solution (with N < 100). But I'm struggling to find a solution with N up to 100.000

Answer 1

To work out user202729's hint:

1. Find the centroid (vertex whose removal leaves subtrees that each have at most half of the vertices of the whole tree).

2. Count the pairs (u, v) whose path contains the centroid.

3. Delete the centroid and operate recursively on each of the subtrees.

Step 1 is O(n) (there's a standard algorithm) and Step 2 will be O(n log n), so the Master Theorem gives O(n log² n) for the whole.

To implement Step 2, root the tree at the centroid and calculate, for the centroid and each of its children, the number of descendants at each depth. Put the results in Fenwick trees for O(log n)-time prefix sums.

Now, sort the edges by weight. Starting with the heaviest edge e and working our way to the lightest, we count the pairs whose path has e as its heaviest edge. The edge e connects a parent and a child; let x be the child. For each (improper, so including x) descendant u of x, let ℓ* = w − K be the greatest ℓ such that w − ℓ ≥ K. With two prefix sums, count the number of vertices v in other subtrees at depth at most ℓ* − depth(u). Then issue updates to the Fenwick trees to remove u from the count.