python graph +广度优先搜索

Posted 2021-05-09

#coding:UTF-8

class Graph:

    def __init__(self):
        # use dict to store the graph
        # the number of keys in self.ves is the number of vertices in the graph
        # the total number of values(list) is the sum of degree of each vertex
        self.ves = {}

    def add_vertex_edges(self, ve):
        v = ve[0]
        neighbours = ve[1]

        if v in self.ves:
            raise KeyError("{} already exists".format(v))
        else:
            self.ves[v] = neighbours

    def neighbours(self, v):
        return self.ves[v]

    def walk(self, start):
        to_visit = [start]
        parent = {start: None}
        cur_level = 0
        level = {start: cur_level}
        print("node:{} level:{} parent:{}".format(start, cur_level, None)) 
        while to_visit:
            next_to_visit = []
            cur_level += 1
            for v in to_visit:
                for n in self.neighbours(v):
                    if n not in level:
                        next_to_visit.append(n)
                        level[n] = cur_level
                        parent[n] = v
                        print("node:{} level:{} parent:{}".format(n, cur_level, v)) 
            to_visit = next_to_visit

    def find_shortest_path(self, v1, v2):
        # breath first search
        class FoundBreak(Exception): pass
        start = v1
        to_visit = [start]
        parent = {start: None}
        cur_level = 0
        level = {start: cur_level}
        try:
            while to_visit:
                next_to_visit = []
                cur_level += 1
                for v in to_visit:
                    for n in self.neighbours(v):
                        if n not in level:
                            next_to_visit.append(n)
                            level[n] = cur_level
                            parent[n] = v
                            if n is v2:
                                raise FoundBreak("")
                to_visit = next_to_visit
        except FoundBreak:
            pass

        # print the path
        path = [v2]
        cur = v2
        while True:
            p = parent[cur]
            path.append(p)
            cur = p
            if p is v1:
                break
        print("find a path from {} to {}".format(v1, v2))
        print(" -> ".join(str(n) for n in reversed(path)))


class Node:

    def __init__(self, s):
        self.s = s

    def __str__(self):
        return self.s

g = Graph()
s = Node("s")
a = Node("a")
z = Node("z")
x = Node("x")
d = Node("d")
c = Node("c")
f = Node("f")
v = Node("v")
g.add_vertex_edges([z, [a]])
g.add_vertex_edges([a, [s, z]])
g.add_vertex_edges([s, [x, a]])
g.add_vertex_edges([x, [s, d, c]])
g.add_vertex_edges([d, [x, c, f]])
g.add_vertex_edges([c, [x, d, f, v]])
g.add_vertex_edges([f, [d, v, c]])
g.add_vertex_edges([v, [f, c]])

g.walk(s)
print("="*30)
g.shortest_path(a, f)