Нахождение соседей вершины в графе - модифицированные BFS

Question

Я новичок в python, поэтому буду очень благодарен за любую помощь. :)

Мне нужно написать функцию, которая возвращает всех соседей данной вершины на заданном расстоянии, используя модифицированные bfs. Мой код указан ниже, но я не знаю, как его исправить.

def bfs_neighbors(graph: List, first_vertex: int, distance: int) -> Set:
    visited, queue = set(), [first_vertex]
    for i in range(distance):
        while queue:
            vertex = queue.pop(0)
            if vertex not in queue:
                visited.add(vertex)
                if not visited:
                    queue.extend(graph[vertex])
        return visited