Visited set

def init_graph():
    # Return adjacency matrix
    return [[0, 10, 6, 5], [10, 0, 0, 15], [6, 0, 0, 4], [5, 15, 4, 0]]
def min_key(keys, visited):
    # Find unvisited index with smallest keys
    best = float('inf')
    idx = -1
    for i in range(len(keys)):
        if not visited[i] and keys[i] < best:
            best = keys[i]
            idx = i
    return idx
def prim(graph):
    # Find MST weight
    n = len(graph)
    keys = [float('inf')] * n
    keys[0] = 0
    visited = [False] * n
    total = 0
    for _ in range(n):
        u = min_key(keys, visited)
        visited[u] = True
        total += keys[u]
        for v in range(n):
            if graph[u][v] and not visited[v] and graph[u][v] < keys[v]:
                keys[v] = graph[u][v]
    return total
print(prim(init_graph()))