Refactor layered_topological_sort method in Graph class

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-01-25 11:43:27 -03:00
commit 0fa85687f1

View file

@ -1,5 +1,5 @@
from typing import Dict, Generator, List, Type, Union
from collections import defaultdict, deque from collections import defaultdict, deque
from typing import Dict, Generator, List, Type, Union
from langchain.chains.base import Chain from langchain.chains.base import Chain
from loguru import logger from loguru import logger
@ -256,7 +256,7 @@ class Graph:
return f"Graph:\nNodes: {vertex_ids}\nConnections:\n{edges_repr}" return f"Graph:\nNodes: {vertex_ids}\nConnections:\n{edges_repr}"
def layered_topological_sort(self): def layered_topological_sort(self):
in_degree = {vertex: 0 for vertex in self.vertices} # Initialize in-degrees in_degree = {vertex.id: 0 for vertex in self.vertices} # Initialize in-degrees
graph = defaultdict(list) # Adjacency list representation graph = defaultdict(list) # Adjacency list representation
# Build graph and compute in-degrees # Build graph and compute in-degrees
@ -265,7 +265,7 @@ class Graph:
in_degree[edge.target_id] += 1 in_degree[edge.target_id] += 1
# Queue for vertices with no incoming edges # Queue for vertices with no incoming edges
queue = deque(vertex for vertex in self.vertices if in_degree[vertex] == 0) queue = deque(vertex.id for vertex in self.vertices if in_degree[vertex.id] == 0)
layers = [] layers = []
current_layer = 0 current_layer = 0
@ -273,9 +273,9 @@ class Graph:
layers.append([]) # Start a new layer layers.append([]) # Start a new layer
layer_size = len(queue) layer_size = len(queue)
for _ in range(layer_size): for _ in range(layer_size):
vertex = queue.popleft() vertex_id = queue.popleft()
layers[current_layer].append(vertex.id) layers[current_layer].append(vertex_id)
for neighbor in graph[vertex]: for neighbor in graph[vertex_id]:
in_degree[neighbor] -= 1 # 'remove' edge in_degree[neighbor] -= 1 # 'remove' edge
if in_degree[neighbor] == 0: if in_degree[neighbor] == 0:
queue.append(neighbor) queue.append(neighbor)