函数式广搜

bfs.py

def expand_dict(record, new_nodes, dist):
    '''record is a dict showes previous distance
    now_nodes is a set for the nodes we found
    dist is a number, showes the distance for new_nodes
    return a dict(record1)
    '''
    new_dict = {i:dist for i in new_nodes}
    #more beautiful assert needed
    assert(set(new_dict.keys()).intersection(set(record.keys())) == set())
    new_record =  {a:b for a,b in list(new_dict.items()) + list(record.items())}
    return new_record


def _bfs(graph, queue, record):
    '''queue is a list
    record is a dict
    return a tuple
    '''
    if queue == None or len(queue) == 0:
        return (None, record)

    assert(set(queue).issubset(set(record.keys())))
    try:
        head = queue[0]
        dist = record[head] + 1
        new_nodes = {x for x in graph.adjacency(head)
                if x not in queue and x not in record.keys()}
    except KeyError:
        print('KeyError:')
        print(queue)
        print(record)
        exit()

    record1 = expand_dict(record, new_nodes, dist)
    queue1 = queue[1:] + list(new_nodes)

    return _bfs(graph, queue1, record1)

def bfs(g):
    '''g is a graph
    Return a dict showes the mimium distance'''
    q0 = [1]
    r0 = dict()
    r0[1] = 0

    (q1, r1) = _bfs(g, q0, r0)
    assert(q1 == None)
    return r1

#sorrow17's code
class Graph:
	def __init__(self):
		self.graph = {}
	def add_edge(self,u,v):
		if u in self.graph.keys() :
			self.graph[u].append(v)
		else:
			self.graph[u] = [v]
	def adjacency(self,u):
		if u in self.graph.keys():
			return self.graph[u]
		else:
			return []

if __name__ == '__main__':
    g = Graph()
    edges = [(1,2),(2,3),(2,5),(1,3),(3,5),(4,5),(2,4),(4,6),(5,6)]
    for u,v in edges:
        g.add_edge(u,v)

    ret = bfs(g)
    print(ret)