main.py

from more_itertools import iterate, locate

def neighbours(a, b):
    raise NotImplementedError # Same as before

def steps_iter(start):
    # infinite iterator with 0-step, 1-step, 2-step, etc. neighbor lists
    return iterate(lambda l: set(sum([neighbours(*p) for p in l], [])), [start])

def solution(start, end):
    # get the ID of the first list which contains `end`
    return next(locate(steps_iter(start), lambda l: end in l))