How do I create a function that helps to find all the subway station within the given distance?

I tried to get previous and next stations within one distance at a time, recursively until distance = 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))

While this may not be the most elegant or efficient code, here is an approach to solving your problem:

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

class Network:
    def __init__(self):
        self._stations = dict()
    
    def add_station(self, sta_name, nxt_sta= None):
        """Adds or updates sta_name and nxt_sta """
        if sta_name in self._stations.keys():
            sta = self._stations[sta_name]
            if nxt_sta:
                sta.add_link(nxt_sta)
        else:
            sta = Station(sta_name, nxt_sta)
            self._stations[sta_name] = sta
        if nxt_sta:
            if nxt_sta in self._stations.keys():
                nxsta = self._stations[nxt_sta]
                nxsta.add_link(sta_name)
            else:
                nxsta = Station(nxt_sta, sta_name)
                self._stations[nxt_sta] = nxsta
    
    @property
    def station(self, sta_name):
        """returns the station node for named station
           throws a key error if station doesn't exist """
        return self._stations[sta_name]
    
    @property
    def size(self):
        """returns size of network in terms of stations"""
        return len(self._stations)
    
    def __repr__(self):
        """return string representation of network"""
        s = ""
        for k, sta in self._stations.items():
            s += f'{sta}\n'           
        return s
        
    def find_all_stations(self, origin, link_count):
        """find list of all stations within link_count of origin
           conforms to breadth first search """
        nodes = set()            #keeps track of stations visited
        stations_found = set()   #keeps track of stations found within range
        level = 0
        sch_list = []            #creates a simple queue for FIFO processing
        if origin in self._stations.keys():
            nodes.add(origin)
            sch_list.append(self._stations[origin])
            while True:
                added_stas = []
                while sch_list:
                    cur_sta = sch_list.pop()
                    sta_lnks = cur_sta.links
                    for lnk in sta_lnks:
                        if lnk not in nodes:
                            added_stas.append(self._stations[lnk])
                if level >= link_count:
                    break
                else:
                    level += 1
                    for s in added_stas:
                        sch_list.append(s)
                        nodes.add(s.name)
                        stations_found.add(s.name)
        return list(stations_found)
    
""" simple class to define data structure for identifying stations 
     add keeping track of neighboring stations"""
class Station:   
    def __init__(self, sta_name, nxt_sta = None):
        self._name = sta_name
        self._lnks = set()
        if nxt_sta:
            self._lnks.add(nxt_sta)   
    @property
    def links(self):
        return list(self._lnks)   
    @property
    def name(self):
        return self._name
    
    def add_link(self, lnk_name):
        self._lnks.add(lnk_name)
        
    def __repr__(self):
        s = f"{self._name} has [ {', '.join(list(self._lnks))} ] neighbors"
        return s
        

Executing the following:

subGrph = Network()
for line in subway_map:
    for s in range(1, len(line)):
        subGrph.add_station(line[s-1], line[s])       
subGrph.find_all_stations('Stevens', 3)        

yields:

['Novena',
 'Dhoby Ghaut',
 'Farrer Park',
 'Toa Payoh',
 'Little India',
 'Newton',
 'Botanic Gardens',
 'Rochor']

def get_stns_dict(mrt_map):
    stations_dict = {}

    for line in mrt_map:
      
        for index, station in enumerate(line):
            # print(station)
            if station not in stations_dict:
                stations_dict[station] = []

            if index == 0: 
                stations_dict[station] += [line[index + 1]]
            elif index == len(line) -1:
                stations_dict[station] += [line[index - 1]]
            else :
                stations_dict[station] += [line[index -1 ]]
                stations_dict[station] += [line[index + 1]]
    return stations_dict

def find_stations_within_distance(mrt_map, orig, dist):
    # Modify the code below
    stations_dict = get_stns_dict(mrt_map)
                
    output = [] 
    possible_stns = stations_dict[orig]  
    past_stns = set()  
    current_distance = 1

    while current_distance <= dist:
        current_distance += 1
        for element in possible_stns:
            past_stns.add(element)
  

        for current_possible_stn in possible_stns:
            next_possible_stns = stations_dict[current_possible_stn]

            for next_stn in next_possible_stns:
                if next_stn not in past_stns and next_stn != orig:
                    output.append(next_stn)

        possible_stns = output
        output = []

    return past_stns