line_finder.py

import numpy as np
import pandas as pd
import pickle

data = pd.read_pickle('dataframe.pkl')

print(data.dtypes)

class LineFinder:

    def __init__(self):
        self.data = data

    def get_the_lines(self,city):
        """
        get all the lines serving a given station

        Parameters:
        station (str): The name of the station.

        Returns:
        lines_of_city (array): The line numbers of the station.
        """
        # Ensure no duplicates in each (Line, City) pair to prevent erroneous data issues
        self.data = self.data.drop_duplicates(subset=['Line', 'Stop name'])
        # Group by 'City' and list the unique 'Line's passing through each city
        lines_per_city = self.data.groupby('Stop name')['Line'].unique()
        lines_of_city= set(lines_per_city.get(city, []))
        return lines_of_city

    def find_common_lines(self, city1 , city2):
        """
        Find common lines between two stations

        Parameters:
        city1, city2 (both str): The name of the stations.

        Returns:
        common_lines (array): The common line numbers between the two stations.
        """
        lines_of_city1 = self.get_the_lines(self ,city1)
        lines_of_city2 = self.get_the_lines(self , city2)
        # Find the intersection (common lines) between the two sets
        common_lines = lines_of_city1.intersection(lines_of_city2)
        return list(common_lines)