Fix BFS depth control in NetworkX graph traversal

This commit is contained in:
yangdx
2025-04-02 23:59:19 +08:00
parent 8f46ce75b8
commit 6d942da338

View File

@@ -303,18 +303,20 @@ class NetworkXStorage(BaseGraphStorage):
# Use BFS to get nodes # Use BFS to get nodes
bfs_nodes = [] bfs_nodes = []
visited = set() visited = set()
queue = [node_label] queue = [(node_label, 0)] # (node, depth) tuple
# Breadth-first search # Breadth-first search
while queue and len(bfs_nodes) < max_nodes: while queue and len(bfs_nodes) < max_nodes:
current = queue.pop(0) current, depth = queue.pop(0)
if current not in visited: if current not in visited:
visited.add(current) visited.add(current)
bfs_nodes.append(current) bfs_nodes.append(current)
# Add neighbor nodes to queue # Only explore neighbors if we haven't reached max_depth
if depth < max_depth:
# Add neighbor nodes to queue with incremented depth
neighbors = list(graph.neighbors(current)) neighbors = list(graph.neighbors(current))
queue.extend([n for n in neighbors if n not in visited]) queue.extend([(n, depth + 1) for n in neighbors if n not in visited])
# Check if graph is truncated - if we still have nodes in the queue # Check if graph is truncated - if we still have nodes in the queue
# and we've reached max_nodes, then the graph is truncated # and we've reached max_nodes, then the graph is truncated