Give minimum permutation weight for edges such that a given set of edge is the Minimum Spanning Tree

Question:

Given a graph of N nodes and M edges, the edges are indexed from 1 -> M. It is guaranteed that there's a path between any 2 nodes.

You need to assign weights for M edges. The weights are in the range of [1...M], and each number can only occur once.

To be shorted, the answer should be a permutation array of [1...M], in which arr[i] = x means edge[i] has the weight of x.

You are given a set R of n-1 edges. R is guaranteed to be a Spanning Tree of the graph.

Find a way to assign weights so that R is the Minimum Spanning Tree of the graph, if there are multiple answers, print the one with minimum lexicographical order.

Contraints:

N, M <= 10^6

Example:

Edges:

3 4
1 2
2 3
1 3
1 4

R = [2, 4, 5]

Answer: 3 4 5 1 2

Explaination:

If you assign weights for the graph like the above image, the MST would be the set R, and it has the smallest lexicographical order.

My take with O(N^2):

Since it asks for the minimum lexicographical order, I traverse through the list of edges, assigning the weights in an increasing order. Intially, w = 1. There can be 3 situations:

• If edge[i] is in R, assign weight[i] = w, increase w by 1
• If edge[i] is not in R: say edge[i] connect nodes u and v. assign weight and increase w for each edge in the path from u to v in R (if that edge is not assigned yet). Then assign weight and increase w for edge[i]
• If edge[i] is assigned, skip it

Is there any way to improve my solution so that it can work in O(N.logN) or less?

Yes, there's an O(m log m)-time algorithm.

The fundamental cycle of a non-tree edge e is comprised of e and the path in the tree between the endpoints of e. Given weights, the spanning tree is minimum if and only if, for every non-tree edge e, the heaviest edge in the fundamental cycle of e is e itself.

The lexicographic objective lends itself to a greedy algorithm, where we find the least valid assignment for edge 1, then edge 2 given edge 1, then edge 3 given the previous edges, etc. Here's the core idea: if the next unassigned edge is a non-tree edge, assign the next numbers to the unassigned tree edges in its fundamental cycle; then assign the next number.

In the example, edge 3-4 is first, and edges 1-3 and 1-4 complete its fundamental cycle. Therefore we assign 1-3 → 1 and 1-4 → 2 and then 3-4 → 3. Next is 1-2, a tree edge, so 1-2 → 4. Finally, 2-3 → 5 (1-2 and 1-3 are already assigned).

To implement this efficiently, we need two ingredients: a way to enumerate the unassigned edges in a fundamental cycle, and a way to assign numbers. My proposal for the former would be to store the spanning tree with the assigned edges contracted. We don't need anything fancy; start by rooting the spanning tree somewhere and running depth-first search to record parent pointers and depths. The fundamental cycle of e will be given by the paths to the least common ancestor of the endpoints of e. To do the contraction, we add a Boolean field indicating whether the parent edge is contracted, then use the path compression trick from disjoint-set forests. The work will be O(m log m) worst case, but O(m) average case. I think there's a strong possibility that the offline least common ancestor algorithms can be plugged in here to get the worst case down to O(m).

As for number assignment, we can handle this in linear time. For each edge, record the index of the edge that caused it to be assigned. At the end, stably bucket sort by this index, breaking ties by putting tree edges before non-tree. This can be done in O(m) time.