Search code examples
pythonmatrixbellman-ford

Find the number of unique nodes in a graph in Python


Given a graph with the format (source, destination, weight):

graph = [
    (1, 2, 2),
    (1, 3, 7),
    (2, 3, 1),
]

I need to assign num_nodes to be the number of unique nodes in the given graph. This is for a Bellman Ford algorithm.

The code I have so far (that is relevant to this question) is as such:

def bellman_ford(graph, source):
    # Find the number of unique nodes in the graph
    # Start
    num_nodes = 0
    uniquevalues = list(set(i for j in graph for i in j))
    for item in uniquevalues:
        num_nodes = num_nodes + 1
    # End

    # Initialize distance vector with infinity and set distance to self is 0
    # Start
    distance = {i : float("Inf") for i in uniquevalues}
    distance[source] = 0
    # End

In my attempt it lists all the unique values of the graph, including the weight values. This obviously causes problem with the output, where it thinks the weight of an edge is a node.

I also had to create uniquevalues to account for initializing the distance vector needing something iterable to assign nodes their distance as infinity. I'm open to suggestions if there are better ways to go about that.


Solution

  • Since it is known that the last item in each tuple in the input list of a graph is not a node, you can unpack the tuple and discard the last item in each iteration before you build a new set. Use the len function to obtain the size of the set for your desired count of unique nodes:

    graph = [
        (1, 2, 2),
        (1, 3, 7),
        (2, 3, 1),
    ]
    unique_nodes = {node for *nodes, _ in graph for node in nodes}
    num_nodes = len(unique_nodes) # num_nodes becomes 3