61 lines
1.8 KiB
Python
61 lines
1.8 KiB
Python
# see the overview from 2026.07.31
|
|
|
|
sample_graph = {
|
|
'Alice': set(['Bob', 'Carol']),
|
|
'Bob': set(['Alice', 'Dave']),
|
|
'Carol': set(['Alice', 'Eve', 'Mallory']),
|
|
'Dave': set(['Frank', 'Eve', 'Bob']),
|
|
'Eve': set(['Carol', 'Grace', 'Dave']),
|
|
'Frank': set(['Dave', 'Heidi']),
|
|
'Grace': set(['Eve', 'Ivan', 'Mallory']),
|
|
'Heidi': set(['Frank', 'Ivan']),
|
|
'Ivan': set(['Heidi', 'Grace']),
|
|
'Mallory': set(['Carol', 'Grace'])
|
|
}
|
|
|
|
|
|
# SOLUTION
|
|
|
|
def nth_degree_connections(
|
|
graph: dict[str, set[str]],
|
|
starting_vertex: str,
|
|
max_n: int
|
|
):
|
|
# what traversal algorithm is this?
|
|
|
|
# since we "visit" when looping through neighbors,
|
|
# mark the starting vertex as visited already
|
|
visited = {starting_vertex}
|
|
queue = [(starting_vertex, 0)] # (vertex, n)
|
|
|
|
to_return: list[list[str]] = []
|
|
|
|
# check for both conditions at once
|
|
# queue is not empty, and we have not exceeded the max_n
|
|
while queue and queue[0][1] <= max_n:
|
|
element = queue.pop(0)
|
|
vertex = element[0]
|
|
current_n = element[1]
|
|
|
|
if current_n == len(to_return):
|
|
to_return.append([])
|
|
to_return[current_n].append(element[0])
|
|
|
|
for neighbor in graph[vertex]:
|
|
if neighbor not in visited:
|
|
queue.append((neighbor, current_n + 1))
|
|
visited.add(neighbor)
|
|
|
|
# we can also
|
|
|
|
# match the example: don't include the 0th-degree connection
|
|
to_return.pop(0)
|
|
return to_return
|
|
|
|
|
|
print(nth_degree_connections(sample_graph, 'Mallory', 2))
|
|
print(nth_degree_connections(sample_graph, 'Alice', 4))
|
|
print(nth_degree_connections(sample_graph, 'Ivan', 1))
|
|
print(nth_degree_connections(sample_graph, 'Grace', 0))
|
|
print(nth_degree_connections(sample_graph, 'Carol', 100))
|