PrivacyScanData/visualize/Visualization-of-popular-algorithms-in-Python/DFS/dfs_visual.py

import networkx as nx
import matplotlib.pyplot as plt
 


#utility fucntion used by DFS which does recursive depth first search 
def DFSUtil(G, v, visited, sl): 
	visited[v] = True
	sl.append(v) 
	for i in G[v]:
		if visited[i] == False:
			DFSUtil(G, i, visited, sl)
	return sl
 


#DFS traversal 
def DFS(G, source): 
	visited = [False]*(len(G.nodes()))
	sl = []		#a list that stores dfs forest starting with source node
	dfs_stk = [] #A nested list that stores all the DFS Forest's
	dfs_stk.append(DFSUtil(G, source, visited, sl))
	for i in range(len(G.nodes())):
		if visited[i] == False:
			sl = []
			dfs_stk.append(DFSUtil(G, i, visited, sl))
	return dfs_stk
			


#takes input from the file and creates a weighted graph
def CreateGraph():
	G = nx.DiGraph()
	f = open('input.txt')
	n = int(f.readline())
	wtMatrix = []
	for i in range(n):
		list1 = map(int,(f.readline()).split())
		wtMatrix.append(list1)
	source = int(f.readline()) #source vertex from where DFS has to start
	#Adds egdes along with their weights to the graph 
	for i in range(n):
		for j in range(n):
			if wtMatrix[i][j] > 0:
					G.add_edge(i, j, length = wtMatrix[i][j]) 
	return G,source



#marks all edges traversed through DFS with red
def DrawDFSPath(G, dfs_stk):
	pos = nx.spring_layout(G)
	nx.draw(G, pos, with_labels = True)  #with_labels=true is to show the node number in the output graph
	edge_labels = dict([((u,v,), d['length']) for u, v, d in G.edges(data = True)])
	nx.draw_networkx_edge_labels(G, pos, edge_labels = edge_labels, label_pos = 0.3, font_size = 11) #prints weight on all the edges
	for i in dfs_stk:
		#if there is more than one node in the dfs-forest, then print the corresponding edges
		if len(i) > 1:
			for j in i[ :(len(i)-1)]:
				if i[i.index(j)+1] in G[j]:
					nx.draw_networkx_edges(G, pos, edgelist = [(j,i[i.index(j)+1])], width = 2.5, alpha = 0.6, edge_color = 'r')
				else:
					#if in case the path was reversed because all the possible neighbours were visited, we need to find the adj node to it.
					for k in i[1::-1]: 
						if k in G[j]:
							nx.draw_networkx_edges(G, pos, edgelist = [(j,k)], width = 2.5, alpha = 0.6, edge_color = 'r')
							break



#main function
if __name__ == "__main__":
	G, source = CreateGraph()
	dfs_stk = DFS(G, source)
	DrawDFSPath(G, dfs_stk)
	plt.show()