python-3.x - 如何创建一个功能来帮助找到给定距离内的所有地铁站？

Question

地铁图如下。它由 3 个不同的列表组成，每个列表代表一条不同的地铁线路。但是，有 2 个连接站可以连接到另一条线路。车站是牛顿站和小印度站。

subway_map = [ ['Botanic Gardens', 'Stevens', 'Newton', 'Little India', 'Rochor'], ['Newton', 'Novena', 'Toa Payoh', 'Braddell', 'Bishan'], ['Dhoby Ghaut', 'Little India', 'Farrer Park', 'Boon Keng']]

所以我想实现一个功能，找到给定距离内的所有站点。例如:

eg_1 = find_station_within_distance(subway_map, origin = 'Botanic Gardens', dist = 3)
eg_2 = find_station_within_distance(subway_map, origin = 'Little India', dist = 1)
eg_3 = find_station_within_distance(subway_map, origin = 'Dhoby Ghaut', dist = 3) 

#function should return either list below

print(eg_1) 
['Stevens', 'Newton'] 
#or 
['Newton', 'Stevens']

print(eg_2)
['Farrer Park', 'Newton', 'Rochor', 'Dhoby Ghaut']

print(eg_3)
['Little India', 'Farrer Park', 'Boon Keng', 'Rochor', 'Newton', 'Stevens', 'Novena']

到目前为止，我只能得到原点所在线路上给定距离内的站点。

def find_stations_within_distance(subway_map, orig, dist):

    result = []

    for lines in subway_map:
    
        if orig in lines:
        
            orig_idx = lines.index(orig)
            max_idx = dist + orig_idx
        
                if max_idx >= len(lines):
            
                    result += lines[orig_idx+1:]
            
                elif max_idx < len(lines):
            
                    result += lines[orig_idx+1:max_idx+1]
            
    return result 

Answer 1

我试图一次在一个距离内获取上一个和下一个站点，递归地直到距离 = 0。

subway_map = [
    ['Botanic Gardens', 'Stevens', 'Newton', 'Little India', 'Rochor'],
    ['Newton', 'Novena', 'Toa Payoh', 'Braddell', 'Bishan'],
    ['Dhoby Ghaut', 'Little India', 'Farrer Park', 'Boon Keng']
]

def find_stations_within_distance(subway_map, orig, dist):
    result = []
    to_find_next = [orig]
    while dist > 0:
        now_finding = to_find_next[:]
        to_find_next.clear()
        while now_finding:
            current_station = now_finding.pop()
            for lines in subway_map:
                if current_station in lines:
                    current_idx = lines.index(current_station)
                    pre_idx = current_idx - 1
                    next_idx = current_idx + 1
                    if pre_idx >= 0:
                        pre_station = lines[pre_idx]
                        if pre_station not in result:
                            to_find_next.append(pre_station)
                            result.append(lines[pre_idx])
                    if next_idx < len(lines):
                        next_station = lines[next_idx]
                        if next_station not in result:
                            to_find_next.append(next_station)
                            result.append(next_station)

        dist -= 1
    if orig in result:
        result.remove(orig)
    return result

# Output: ['Stevens', 'Newton']
print(find_stations_within_distance(subway_map, 'Botanic Gardens', 2))
# Output: ['Newton', 'Rochor', 'Dhoby Ghaut', 'Farrer Park']
print(find_stations_within_distance(subway_map, 'Little India', 1))
# Output: ['Little India', 'Newton', 'Rochor', 'Farrer Park', 'Boon Keng', 'Stevens', 'Novena']
print(find_stations_within_distance(subway_map, 'Dhoby Ghaut', 3))