bfs and dfs for graph

graph.py

#! usr/local/bin python3

# implementation for graph algorithms

class Graph:
	def __init__(self, n):
		# n is the number of nodes
		self.n=n
		# adj is the adjacency list
		self.adj=[[] for _ in range(n)]
	def add_edge(self, u, v):
		self.adj[u].append(v)
		self.adj[v].append(u)
	
	def bfs(self,s):
		"""
		breadth first search
		s: start node
		return: visited nodes
		"""
		visited=[False for _ in range(self.n)]
		queue=[s]
		visited[s]=True
		while queue:
			u=queue.pop(0)
			for v in self.adj[u]:
				if not visited[v]:
					queue.append(v)
					visited[v]=True
		return visited
	
	def dfs(self, s):
		"""
		depth first search
		s:start node
		return: visited nodes
		"""
		visited =[False for _ in range(self.n)]
		stack=[s]
		visited[s]=True
		while stack:
			u=stack.pop()
			for v in self.adj[u]:
				if not visited[v]:
					stack.append(v)
					visited[v]=True
		return visited

	# function for finding the shortest path for the graph
	def shortest_path(self, s, t):
		"""
		shortest path
		s: start node
		t: end node
		return: the shortest path
		"""
		visited=[False for _ in range(self.n)]
		queue=[s]
		visited[s]=True
		prev=[None for _ in range(self.n)]
		while queue:
			u=queue.pop(0)
			for v in self.adj[u]:
				if not visited[v]:
					queue.append(v)
					visited[v]=True
					prev[v]=u
		return self.get_path(prev,s,t)
	def get_path(self,prev,s,t):
		"""
		get the path
		prev:previous node
		s; start node
		t: end node
		return: the path
		"""
		res=[]
		if prev[t] is not None or s==t:
			while t is not None:
				res.insert(0,t)
				t=prev[t]
		return res

if __name__ == '__main__':
	g = Graph(6)
	g.add_edge(0, 1)
	g.add_edge(0, 2)
	g.add_edge(1, 2)
	g.add_edge(1, 3)
	g.add_edge(2, 4)
	g.add_edge(3, 4)
	g.add_edge(3, 5)
	g.add_edge(4, 5)
	print(g.bfs(0))
	print(g.dfs(0))
	print(g.shortest_path(0,5))