# -*- coding: utf-8 -*-
"""Graph data structure with common traversal and path-finding algorithms."""
import json
from queue import Queue
from typing import Any, Dict, List, Optional
from .vertex import Vertex
[docs]
class Graph:
"""Undirected weighted graph backed by an adjacency dictionary of Vertex objects."""
[docs]
def __init__(self):
self.vertices = {}
[docs]
@classmethod
def create_example_graph(cls):
r"""
Example graph structure::
A
/ \
B --- C
\ /
D
/ \
E --- F
"""
return cls.build_graph({
"A": ["B", "C"],
"B": ["C", "D"],
"C": ["D"],
"D": ["E", "F"],
"E": ["F"]
})
def __iter__(self):
return iter(self.vertices.values())
def __contains__(self, key):
return key in self.vertices
def __str__(self):
return json.dumps({
key: list(vertex.get_connections_ids())
for key, vertex in self.vertices.items()
})
def __repr__(self):
return json.dumps({
key: list(vertex.get_connections_ids())
for key, vertex in self.vertices.items()
})
[docs]
@classmethod
def build_graph(cls, graph_definition: Dict[Any, List]):
"""
Create a Graph
:param graph_definition: Graph definition
:return: The Graph.
>>> str(Graph.create_example_graph())
'{"A": ["B", "C"], "B": ["C", "D"], "C": ["D"], "D": ["E", "F"], "E": ["F"], "F": []}'
"""
graph = cls()
for key, vertices in graph_definition.items():
for vertex in vertices:
graph.add_edge(key, vertex)
return graph
[docs]
def add_vertex(self, key) -> Vertex:
""" Add a vertex for the given key if it does not already exist and return it. """
if key not in self.vertices:
self.vertices[key] = Vertex(key)
return self.vertices[key]
[docs]
def get_vertex(self, key) -> Optional[Vertex]:
""" Return a vertex by key """
return self.vertices.get(key, None)
[docs]
def get_vertices(self) -> Dict:
""" Return all vertices """
return self.vertices
[docs]
def add_edge(self, key_vertex1, key_vertex2, weight=0):
""" Add a directed edge from key_vertex1 to key_vertex2 with an optional weight. """
if key_vertex1 not in self.vertices:
self.add_vertex(key_vertex1)
if key_vertex2 not in self.vertices:
self.add_vertex(key_vertex2)
self.vertices[key_vertex1].add_neighbor(self.vertices[key_vertex2], weight)
[docs]
def breadth_first_search(self, node):
"""
Algorithm used for tree traversal on graphs or tree data structures.
:param node: Initial node.
:return: List of visited nodes.
>>> Graph.create_example_graph().breadth_first_search("A")
['A', 'B', 'C', 'D', 'E', 'F']
"""
visited, queue = [node], Queue()
queue.put(node)
while not queue.empty():
queue.get()
for neighbour in self.get_vertices():
if neighbour not in visited:
visited.append(neighbour)
queue.put(neighbour)
return visited
[docs]
def depth_first_search(self, node_id, visited: Optional[List] = None):
"""
Depth First Traversal Algorithm...
:param node_id: Initial node.
:param visited: Set of visited nodes.
:return: List of visited nodes.
>>> Graph.create_example_graph().depth_first_search("A")
['A', 'B', 'C', 'D', 'E', 'F']
"""
if not visited:
visited = []
visited.append(node_id)
for neighbour in self.vertices[node_id].get_connections():
if neighbour.id not in visited:
self.depth_first_search(neighbour.id, visited)
return visited
[docs]
def find_path(self, start, end, path: Optional[List] = None):
"""
Determine a path between two nodes...
:param start: Start node.
:param end: End node.
:param path:
:return: The path.
>>> Graph.create_example_graph().find_path("A", "D")
["A", "B", "C", "D"]
"""
if not path:
path = []
path = path + [start]
if start == end:
return path
if start not in self:
return None
for node in self.vertices[start].get_connections():
if node.id not in path:
new_path = self.find_path(node.id, end, path)
if new_path:
return new_path
return None
[docs]
def find_all_paths(self, start, end, path: Optional[List] = None):
"""
Determine all paths between two nodes...
:param start: Start node.
:param end: End node.
:param path:
:return: The path.
>>> Graph.create_example_graph().find_all_paths("A", "D")
[["A", "B", "C", "D"], ["A", "B", "D"], ["A", "C", "D"]]
"""
if not path:
path = []
path = path + [start]
if start == end:
return [path]
if start not in self:
return []
paths = []
for node in self.vertices[start].get_connections():
if node.id not in path:
new_paths = self.find_all_paths(node.id, end, path)
paths.extend(new_paths)
return paths
[docs]
def find_shortest_path(self, start, end):
"""
Find the shortest path between two nodes...
:param start: Start node.
:param end: End node.
:return: Shortest path.
>>> Graph.create_example_graph().find_shortest_path("A", "F")
['A', 'B', 'D', 'F']
>>> Graph.create_example_graph().find_shortest_path("B", "E")
['B', 'D', 'E']
"""
dist, q = {start: [start]}, Queue()
q.put(start)
while not q.empty():
key = q.get()
for next_ in self.vertices[key].get_connections():
next_ = next_.id
if next_ not in dist:
dist[next_] = dist[key] + [next_]
q.put(next_)
return dist.get(end)