priority first search on a graphviz dot file

pfs.py

#!/usr/bin/python3

# See this article for more explanations:
# https://ckhung.medium.com/4-in-1-priority-first-search-in-python-bfs-dfs-prims-and-dijkstra-s-algorithms-4e9fe8ccba87

# pip3 install pqdict pydot networkx
# python3 pfs.py -a dijk -0 E t02.dot
# python3 pfs.py -a prim -0 G t02.dot

# https://github.com/nvictus/priority-queue-dictionary
# https://pypi.org/project/pydot/
# https://networkx.org/documentation/stable/tutorial.html

import argparse, sys, pydot, re
import networkx as nx
from pqdict import pqdict

parser = argparse.ArgumentParser(
    description=u'BFS/DFS/PFS all in one',
    formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('-0', '--start', type=str, default='', 
    help='starting vertex')
algos = ['bfs', 'dfs', 'prim', 'dijk']
parser.add_argument('-a', '--algo', type=str, default='', 
    help=f'type of algorithm, one of {algos}')
parser.add_argument('graph_file', help='graphviz .dot file')
args = parser.parse_args()

if not args.algo in algos:
    print(f'must specify -a ... and it must be one of {algos}')
    sys.exit()
G = pydot.graph_from_dot_file(args.graph_file)[0]
G = nx.nx_pydot.from_pydot(G)

#for v in G.nodes:
#    G.nodes[v]['#priority'] = float("inf")
#    print(v, G.nodes[v], G.adj[v])

fringe = pqdict(key=lambda vdata: vdata['#priority'])
if not args.start:
    args.start = sorted(G.nodes)[0]

G.nodes[args.start]['#priority'] = 0
fringe.additem(args.start, G.nodes[args.start])
vid = 0
while fringe:
    (vname, vdata) = fringe.popitem()
    vdata['#status'] = 'visited'
    print('{}({}) | '.format(vname, vdata['#priority']), end=' ')
    curr_name = vname
    curr_data = vdata
    while '#parent' in G.nodes[curr_name]:
        curr_name = curr_data['#parent']
        curr_data = G.nodes[curr_name]
        print(curr_name, end=' ')
    print('')
    for wname in sorted(G.adj[vname].keys()):
        wdata = G.nodes[wname]
        edge = G.adj[vname][wname][0]
        edge['#weight'] = int(re.sub(r'[^\d\.]', '', edge['label']))
        # print('{} ({})'.format(wname, edge['#weight']), end=' ')
        if not '#status' in wdata: vid += 1
        if args.algo == 'bfs':
            new_prio = vid
        elif args.algo == 'dfs':
            new_prio = -vid
        elif args.algo == 'prim':
            new_prio = edge['#weight']
        elif args.algo == 'dijk':
            new_prio = vdata['#priority'] + edge['#weight']
        if not '#status' in wdata:
            wdata['#status'] = 'fringe'
            wdata['#parent'] = vname
            wdata['#priority'] = new_prio
            fringe.additem(wname, wdata)
        elif wdata['#status'] == 'fringe' and args.algo in ['prim', 'dijk']:
            if new_prio < wdata['#priority']:
                wdata['#parent'] = vname
                wdata['#priority'] = new_prio
                fringe.updateitem(wname, wdata)
                fringe.heapify(wname)

if args.algo == 'prim':
    total = 0
    for vname in G.nodes:
        vdata = G.nodes[vname]
        if '#parent' in vdata:
            total += G[vname][vdata['#parent']][0]['#weight']
    print('MST total cost: {}'.format(total))

t01.dot

graph G {
    rankdir = LR;

    A -- B [ label="9" ];
    A -- C [ label="4" ];
    A -- D [ label="3" ];
    A -- E [ label="8" ];
    B -- C [ label="5" ];
    B -- D [ label="4" ];
    B -- E [ label="2" ];
    C -- D [ label="7" ];
    D -- E [ label="6" ];
}

t02.dot

# https://www.geeksforgeeks.org/dijkstras-shortest-path-algorithm-using-priority_queue-stl/
# https://www.geeksforgeeks.org/prims-minimum-spanning-tree-mst-greedy-algo-5/
# Vertex "1" in geeksforgeeks is vertex "A" here, "2" is "B", "3" is "C",
# ... "8" is "H", and vertex "0" there is named "I" here,

graph G {
    rankdir = LR;

    A -- B  [ label="8" ];
    A -- G  [ label="11" ];
    A -- I  [ label="4" ];
    B -- C  [ label="7" ];
    B -- E  [ label="4" ];
    B -- H  [ label="2" ];
    C -- D  [ label="9" ];
    C -- E  [ label="14" ];
    D -- E  [ label="10" ];
    E -- F  [ label="2" ];
    F -- G  [ label="1" ];
    F -- H  [ label="6" ];
    G -- H  [ label="7" ];
    G -- I  [ label="8" ];
}