cle.py

# Import the required libraries
from collections import defaultdict

# Define the Chu-Liu/Edmonds algorithm function
def chu_liu_edmonds(graph):
    # Step 1: Contract the graph to a single node
    contracted_node = max(graph.keys()) + 1
    contracted_graph = defaultdict(list)
    for node in graph:
        for (weight, child) in graph[node]:
            contracted_graph[node].append((weight, child))
            contracted_graph[contracted_node].append((0, child))
    
    # Step 2: Initialize the arrays for the algorithm
    parent = [None] * (contracted_node + 1)
    max_weight = [float('-inf')] * (contracted_node + 1)
    for node in contracted_graph:
        for (weight, child) in contracted_graph[node]:
            max_weight[child] = max(max_weight[child], weight)
    
    # Step 3: Find the maximum incoming edges for each node
    for node in contracted_graph:
        for (weight, child) in contracted_graph[node]:
            if max_weight[child] == weight:
                parent[child] = node
    
    # Step 4: Find cycles and contract them
    visited = [False] * (contracted_node + 1)
    visited[contracted_node] = True
    cycle = None
    for node in contracted_graph:
        if not visited[node]:
            visited[node] = True
            path = [node]
            while True:
                child = parent[node]
                if visited[child]:
                    cycle = path[path.index(child):]
                    break
                visited[child] = True
                path.append(child)
                node = child
            if cycle is not None:
                break
    if cycle is not None:
        # Step 4(a): Find the minimum weight edge in the cycle
        min_weight = float('inf')
        min_edge = None
        for node in cycle:
            for (weight, child) in contracted_graph[node]:
                if child in cycle and weight < min_weight:
                    min_weight = weight
                    min_edge = (node, child)
        
        # Step 4(b): Contract the cycle by creating a new node
        new_node = max(graph.keys()) + 1
        graph[new_node] = []
        for node in contracted_graph:
            for (weight, child) in contracted_graph[node]:
                if node == min_edge[0] and child == min_edge[1]:
                    graph[node].remove((weight, child))
                    graph[node].append((0, new_node))
                elif child in cycle:
                    graph[node].remove((weight, child))
                    graph[node].append((weight - min_weight, new_node))
        graph[new_node] = [(min_weight, child) for child in cycle if child != min_edge[1]]
        
        # Step 4(c): Recurse on the contracted graph
        chu_liu_edmonds(graph)
    
    # Step 5: Find the MST in the contracted graph
    mst = {}
    for node in graph:
        for (weight, child) in graph[node]:
            if child != contracted_node:
                mst[node] = (weight, child)
    
    return mst