Группировка людей в pandas фрейме данных с индивидуальной функцией

Question

Введение: У меня есть pandas кадр данных с людьми, которые живут в разных местах (широта, долгота, номер этажа). Я хочу объединить 3 человека в каждой группе. Это означает, что в конце этого процесса каждый человек относится к одной конкретной группе. Мой фрейм данных имеет длину, кратную 9 (например, 18 человек).

Часть tricky состоит в том, что людям в одной группе не разрешается находиться в одном и том же месте с точки зрения широты и долготы.

Что происходит не так? После того, как я применил свою функцию к pandas фрейму данных, я получил новый фрейм данных, в котором люди распределены по группам. Однако есть 2 проблемы:

1) Каждый раз, когда 3 человека не назначаются группе! Я не знаю почему, и я не знаю, является ли это причиной моей второй проблемы.

2) По некоторым причинам, которые мне не ясны, люди с одинаковым местоположением (широта / долгота) сгруппированы вместе! Помните, что это никогда не должно произойти.

Вот что я сделал:

(ссылка на Google Colab )

Импорт библиотека:

from math import atan2, cos, radians, sin, sqrt
import math
import pandas as pd
import numpy as np
from random import random

Получить данные:

array_data=([[ 50.56419  ,   8.67667  ,   2.       , 160.       ],
   [ 50.5740356,   8.6718179,   1.       ,   5.       ],
   [ 50.5746321,   8.6831284,   3.       , 202.       ],
   [ 50.5747453,   8.6765588,   4.       , 119.       ],
   [ 50.5748992,   8.6611471,   2.       , 260.       ],
   [ 50.5748992,   8.6611471,   3.       , 102.       ],
   [ 50.575    ,   8.65985  ,   2.       , 267.       ],
   [ 50.5751   ,   8.66027  ,   2.       ,   7.       ],
   [ 50.5751   ,   8.66027  ,   2.       ,  56.       ],
   [ 50.57536  ,   8.67741  ,   1.       , 194.       ],
   [ 50.57536  ,   8.67741  ,   1.       , 282.       ],
   [ 50.5755255,   8.6884584,   0.       , 276.       ],
   [ 50.5755273,   8.674282 ,   3.       , 167.       ],
   [ 50.57553  ,   8.6826   ,   2.       , 273.       ],
   [ 50.5755973,   8.6847492,   0.       , 168.       ],
   [ 50.5756757,   8.6846139,   4.       , 255.       ],
   [ 50.57572  ,   8.65965  ,   0.       ,  66.       ],
   [ 50.57591  ,   8.68175  ,   1.       , 187.       ]])

all_persons = pd.DataFrame(data=array_data) # convert back to dataframe

all_persons.rename(columns={0: 'latitude', 1: 'longitude', 2:'floor', 3:'id'}, inplace=True) # rename columns

Это функция для расчета расстояния между людьми. Если расстояние равно 0, люди имеют одинаковое местоположение с точки зрения широты и долготы.

def calculate_distance(lat1, lon1, lat2, lon2, floor_person_1, floor_person_2):
    """
    Calculate the shortest distance between two points given by the latitude and
    longitude.
    """
    scattering_factor = 0.0001

    earth_radius = 6373  # Approximate / in km.
    lat1 = radians(lat1)
    lon1 = radians(lon1)
    lat2 = radians(lat2)
    lon2 = radians(lon2)

    dlon = lon2 - lon1
    dlat = lat2 - lat1

    a = sin(dlat / 2) ** 2 + cos(lat1) * cos(lat2) * sin(dlon / 2) ** 2
    c = 2 * atan2(sqrt(a), sqrt(1 - a))

    distance = earth_radius * c  # Unit: km. Parameter does not consider floor number
    print(distance)

    # people share same location (latitude / longitude) but on different floors 
    # --> its ok to put them in same group
    if (distance==0) and (floor_person_1 != floor_person_2): 
      distance = distance + scattering_factor
      print('Identical location but different floors')
      print('lat1:', math.degrees(lat1), 'lon1:', math.degrees(lon1))
      print('lat2', math.degrees(lat2), 'lon2:', math.degrees(lon2)) 

    else: # people share different locations (latitude / longitude)
      pass

    return distance    

Это моя функция, которая группирует людей:

def group_people(all_persons, max_distance_parameter):

    assert len(all_persons) % 9 == 0
    all_persons.set_index("id", drop=False, inplace=True)

    all_persons["host"] = np.nan
    all_persons["group"] = np.nan

    Streufaktor = 0.0001 
    max_distance = max_distance_parameter
    group_number = 0
    group = []
    for index, candidate in all_persons.iterrows():
        if len(group) == 3:
            for person in group:
                all_persons.at[person["id"], "group"] = group_number
            group_number += 1
            group = []

        if len(group) == 0:
            group.append(candidate)
        else:
            for person in group:
                distance = calculate_distance(
                    candidate["latitude"],
                    candidate["longitude"],
                    person["latitude"],
                    person["longitude"],
                    candidate['floor'],
                    person['floor']
                )
                distance = distance 

                if 0 < distance <= max_distance:
                    group.append(candidate)
                    break

Затем я применяю функцию к фрейм данных и посмотрите на результат:

group_people(all_persons,4)

all_persons

Вот что я получаю:

Желтым цветом вы видите, что пошло не так (см. Определение проблемы выше).

Как я могу это исправить? (пожалуйста, проверьте связанный Google Colab)