You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
122 lines
3.8 KiB
122 lines
3.8 KiB
"""
|
|
:Author: James Sherratt
|
|
:Date: 20/10/2019
|
|
:License: MIT
|
|
|
|
:name: graph.py
|
|
"""
|
|
|
|
|
|
class Graph:
|
|
|
|
def __init__(self):
|
|
self.vertices = set()
|
|
self.edges = set()
|
|
|
|
def add_vertex(self, vert):
|
|
"""
|
|
Add a vertex to the graph.
|
|
|
|
:param vert: name of the vertex.
|
|
:return: None
|
|
"""
|
|
self.vertices.add(vert)
|
|
|
|
def add_edge(self, vert1, vert2, directional=False):
|
|
"""
|
|
Add an edge to the graph. The edge will be defined as a simple tuple where:
|
|
- The first value is the initial edge.
|
|
- The second is the final edge.
|
|
- The third is whether the edge is directional.
|
|
|
|
:param vert1: the start vertex.
|
|
:param vert2: the end vertex.
|
|
:param directional: whether or not the edge has a direction. Default: False
|
|
:return: None
|
|
"""
|
|
self.vertices.add(vert1)
|
|
self.vertices.add(vert2)
|
|
if (not directional) and (vert1 > vert2):
|
|
# swap if not directional to avoid duplicates in the set.
|
|
self.edges.add((vert2, vert1, directional))
|
|
else:
|
|
self.edges.add((vert1, vert2, directional))
|
|
|
|
def adjacency(self, vert1, vert2):
|
|
"""
|
|
Check if 2 vertices are adjacent (if they exist). Note: if vert1 and vert2
|
|
are connected, but directionally from vert2 to vert1, False is returned.
|
|
|
|
:param vert1: The first vertex to compare.
|
|
:param vert2: The second vertex to compare.
|
|
:return: True if adjacent, False if not, None if either are not in the set.
|
|
"""
|
|
if (vert1 not in self.vertices) or (vert2 not in self.vertices):
|
|
return None
|
|
|
|
for edge in self.edges:
|
|
if (vert1 == edge[0]) and (vert2 == edge[1]):
|
|
return True
|
|
if (not edge[2]) and ((vert1 == edge[1]) and (vert2 == edge[0])):
|
|
return True
|
|
return False
|
|
|
|
def neighbours(self, vert):
|
|
"""
|
|
Get the neighbours of a vertex.
|
|
:param vert: name of the vertex.
|
|
:return: list of neighbours, or None if the vertex is not in the graph.
|
|
"""
|
|
if vert not in self.vertices:
|
|
return None
|
|
|
|
neighbours = set()
|
|
for edge in self.edges:
|
|
if vert == edge[0]:
|
|
neighbours.add(edge[1])
|
|
elif (not edge[2]) and (vert == edge[1]):
|
|
neighbours.add(edge[0])
|
|
return neighbours
|
|
|
|
def as_dict(self):
|
|
"""
|
|
Convert the graph to a dictionary where:
|
|
- Each key is a vertex.
|
|
- Each value is a set of neighbours you can travel to.
|
|
|
|
:return: dict representing the graph.
|
|
"""
|
|
graph_dict = {v: set() for v in self.vertices}
|
|
for edge in self.edges:
|
|
graph_dict[edge[0]].add(edge[1])
|
|
if not edge[2]:
|
|
graph_dict[edge[1]].add(edge[0])
|
|
|
|
return graph_dict
|
|
|
|
|
|
if __name__ == "__main__":
|
|
mygraph = Graph()
|
|
mygraph.add_edge("a", "b")
|
|
mygraph.add_edge("b", "c")
|
|
mygraph.add_edge("b", "d")
|
|
mygraph.add_edge("c", "b")
|
|
mygraph.add_edge("a", "e", directional=True)
|
|
mygraph.add_vertex("z")
|
|
print("b neighbours:", mygraph.neighbours("b"))
|
|
print("a neighbours:", mygraph.neighbours("a"))
|
|
print("q neighbours:", mygraph.neighbours("q"))
|
|
print("e neighbours:", mygraph.neighbours("e"))
|
|
print()
|
|
# Adjacency has direction.
|
|
print("a and e adjacent:", mygraph.adjacency("a", "e"))
|
|
print("e and a adjacent:", mygraph.adjacency("e", "a"))
|
|
print("d and b adjacent", mygraph.adjacency("d", "b"))
|
|
print("q and a adjacent:", mygraph.adjacency("q", "a"))
|
|
print("z and a adjacent:", mygraph.adjacency("z", "a"))
|
|
print()
|
|
print("as dict")
|
|
print(mygraph.as_dict())
|
|
|
|
# Exercise/ project: add a method "path" to find path between 2 vertices.
|
|
# (Hint: lookup DFS/ BFS algorithm.)
|