This code is meant to implement a
Graph class which has a method that returns a path between 2 nodes using breadth-first search.
I'm doing this to improve my style and to improve my knowledge of fundamental algorithms/data structures for an upcoming coding interview.
from queue import PriorityQueue from collections import deque from copy import deepcopy class Graph: def __init__(self): self.nodes = dict() def add_node(self, id): new_node = Node(id) self.nodes[id] = new_node return new_node def get_node(self, id): node = self.nodes[id] if id in self.nodes else None return node def bfs(self, start, end): queue = PriorityQueue() queue.put((0, start)) seen = set([start]) parents = dict() while not queue.empty(): cur_node = queue.get(0) if cur_node is not end: for neighbour in cur_node.neighbours: if neighbour not in seen: dist_to_neighbour = cur_node.neighbours[neighbour] queue.put((dist_to_neighbour, neighbour)) parents[neighbour] = cur_node seen.add(neighbour) else: path = self.get_path(start, end, parents) return path return None def get_path(self, start, end, parents): path = deque([end]) cur_node = end while cur_node is not start: cur_node = parents[cur_node] path.appendleft(cur_node) return list(path) class Node: def __init__(self, id): self.id = id self.neighbours = dict() def add_neighbour(self, to, weight): self.neighbours[to] = weight
I have a few ideas about what might be improved but I'm not sure how to do it (or even if it's worth doing).
- There might be a better way to do
get_node()instead of squashing it into 2 lines like how I've done it.
cur_node = queue.get(0)should be 2 lines?
dist_to_neighbour = cur_node.neighbours[neighbour]to use a method instead of using dictionary notation (which might be a little awkward). Maybe I should have a method instead?
- Using something else instead of a
get_path()(especially in an interview context where deque might be a bit too exotic).
- Have more comments?