Minimize total distance between drivers and their assigned bus station given capacity constraints

bus_park_optimization.py

import pulp
import pandas as pd
import numpy as np

"""Example drivers.csv file:
driver_id,lat,lon
dr_1,10,20
dr_2,30,50
dr_3,50,40
dr_4,80,20
dr_5,40,80
dr_6,60,60
dr_7,70,70
dr_8,90,90
dr_9,100,30
dr_10,40,85
dr_11,60,60
dr_12,70,70"""

"""Example stations.csv file:
station_id,lat,lon,capacity
st_1,10,30,3
st_2,50,50,3
st_3,50,20,3
st_4,20,10,3"""


def optimize_bus_stations(drivers, stations):
    """
    :param stations: pandas.DataFrame with columns 'station_id', 'lat', 'lon', 'capacity'
    :param drivers: pandas.DataFrame with columns 'driver_id', 'lat', 'lon'
    :return: dict that maps driver_id to station_id
    """

    num_drivers = drivers.shape[0]
    num_stations = stations.shape[0]

    assignment = pulp.LpVariable.dicts(
        "Assign",
        ((i, j) for i in range(num_drivers) for j in range(num_stations)),
        cat='Binary'
    )

    distances = np.array(drivers.apply(lambda row: stations.apply(lambda row2: np.sqrt((row.lat-row2.lat)**2 + (row.lon-row2.lon)**2), axis=1), axis=1))

    prob = pulp.LpProblem("Bus_Driver_Parking", pulp.LpMinimize)
    prob += pulp.lpSum(assignment[i, j] * distances[i][j] for i in range(num_drivers) for j in range(num_stations))
    for i in range(num_drivers):
        prob += pulp.lpSum(assignment[i, j] for j in range(num_stations)) == 1
    for j in range(num_stations):
        print(stations.iloc[j])
        prob += pulp.lpSum(assignment[i, j] for i in range(num_drivers)) <= stations.iloc[j].capacity
    prob.solve()

    assignment_dict = {}
    for i in range(num_drivers):
        for j in range(num_stations):
            if assignment[i, j].value() == 1.0:
                assignment_dict[drivers.iloc[i].driver_id] = stations.iloc[j].station_id
    return assignment_dict, prob.objective.value()


drivers = pd.read_csv('drivers.csv')
stations = pd.read_csv('stations.csv')

assignment, objective = optimize_bus_stations(drivers, stations)
print(f'Assignment: {assignment}')
print(f'Total distance: {objective}')