diff --git a/graph_search.py b/graph_search.py index e76d2e39..e13fafa8 100644 --- a/graph_search.py +++ b/graph_search.py @@ -11,53 +11,39 @@ def __init__(self, graph): self.graph = graph def find_path(self, start, end, path=None): - self.start = start - self.end = end - self.path = path if path else [] + path = path or [] - self.path += [self.start] - if self.start == self.end: - return self.path - if self.start not in self.graph: - return None - for node in self.graph[self.start]: - if node not in self.path: - newpath = self.find_path(node, self.end, self.path) + path.append(start) + if start == end: + return path + for node in self.graph.get(start, []): + if node not in path: + newpath = self.find_path(node, end, path) if newpath: return newpath - return None def find_all_path(self, start, end, path=None): - self.start = start - self.end = end - _path = path if path else [] - _path += [self.start] - if self.start == self.end: - return [_path] - if self.start not in self.graph: - return [] + path = path or [] + path.append(start) + if start == end: + return [path] paths = [] - for node in self.graph[self.start]: - if node not in _path: - newpaths = self.find_all_path(node, self.end, _path[:]) - for newpath in newpaths: - paths.append(newpath) + for node in self.graph.get(start, []): + if node not in path: + newpaths = self.find_all_path(node, end, path[:]) + paths.extend(newpaths) return paths def find_shortest_path(self, start, end, path=None): - self.start = start - self.end = end - _path = path if path else [] + path = path or [] + path.append(start) - _path += [self.start] - if self.start == self.end: - return _path - if self.start not in self.graph: - return None + if start == end: + return path shortest = None - for node in self.graph[self.start]: - if node not in _path: - newpath = self.find_shortest_path(node, self.end, _path[:]) + for node in self.graph.get(start, []): + if node not in path: + newpath = self.find_shortest_path(node, end, path[:]) if newpath: if not shortest or len(newpath) < len(shortest): shortest = newpath