Вернуть расстояние до ближайшей точки по группе - python

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

Изменить: добавлен вариант, который (должен) включать идентификатор ближайшей точки.

from sklearn.neighbors import BallTree
import pandas as pd
import numpy as np
from scipy.spatial.distance import pdist, squareform

df = pd.DataFrame({              
    'Time' : [1,1,1,1,1,1,2,2,2,2,2,2],             
    'ID' : ['A','B','C','X','U','V','A','B','C','X','U','V'],      
    'Group' : ['Red','Red','Red','Grn','Grn','Grn','Red','Red','Red','Grn','Grn','Grn'],           
    'X' : [2.0,3.0,4.0,2.0,2.0,1.0,1.0,6.0,4.0,2.0,5.0,3.0],
    'Y' : [3.0,1.0,0.0,0.0,2.0,1.0,2.0,0.0,1.0,1.0,0.0,0.0],           
    })


def subgroup_distance(df, group_column='Group', include_point_itself=True):    
    groups = df[group_column].unique()

    all_points = df[['X','Y']].values
    
    for group in groups:
        group_points = df[df[group_column] == group][['X','Y']]
        tree = BallTree(group_points, leaf_size=15, metric='minkowski')
        
        if include_point_itself:
            distance, index = tree.query(all_points, k=1)
            distances = distance[:,0]
            distance_column_name = "distance_{}".format( group )
            df[ distance_column_name ] = distances
                
        else:
            indici = (df[group_column] == group).values * 1
            distance, index = tree.query(all_points, k=2)
            distances = distance[ np.arange(distance.shape[0]), indici]

            distance_column_name = "distance_{}".format( group )
            df[ distance_column_name ] = distances

    return df

def calculate_distances(df):
    return subgroup_distance(df, include_point_itself=False)

df_distances = df.groupby(['Time']).apply(calculate_distances).reset_index()

это выведет

    index  Time ID Group    X    Y  distance_Red  distance_Grn
0       0     1  A   Red  2.0  3.0      2.236068      1.000000
1       1     1  B   Red  3.0  1.0      1.414214      1.414214
2       2     1  C   Red  4.0  0.0      1.414214      2.000000
3       3     1  X   Grn  2.0  0.0      1.414214      1.414214
4       4     1  U   Grn  2.0  2.0      1.000000      1.414214
5       5     1  V   Grn  1.0  1.0      2.000000      1.414214
6       6     2  A   Red  1.0  2.0      3.162278      1.414214
7       7     2  B   Red  6.0  0.0      2.236068      1.000000
8       8     2  C   Red  4.0  1.0      2.236068      1.414214
9       9     2  X   Grn  2.0  1.0      1.414214      1.414214
10     10     2  U   Grn  5.0  0.0      1.000000      2.000000
11     11     2  V   Grn  3.0  0.0      1.414214      1.414214

Вариант, который выведет ID ближайшей точки в подгруппе

def subgroup_distance_with_nearest(df, group_column='Group', include_point_itself=True):    
    groups = df[group_column].unique()

    all_points = df[['X','Y']].values
    
    for group in groups:
        group_points = df[df[group_column] == group][['X','Y']]
        tree = BallTree(group_points, leaf_size=15, metric='minkowski')
        
        distances = None
        nearest_id = None
        
        if include_point_itself:
            distance, index = tree.query(all_points, k=1)
            distances = distance[:,0]
            nearest_id = group_points.index[index[:,0]]
                        
        else:
            indici = (df[group_column] == group).values * 1
            distance, index = tree.query(all_points, k=2)
            distances = distance[ np.arange(distance.shape[0]), indici]
            index_indici =  index[ np.arange(distance.shape[0]), indici]
            nearest_id = group_points.index[index_indici]

        distance_column_nearest_name = "nearest_index_{}".format( group )
        distance_column_name = "distance_{}".format( group )
        df[ distance_column_name ] = distances
        df[ distance_column_nearest_name] = nearest_id


    return df

def subgroup_distance_with_nearest(df):
    return subgroup_distance(df, include_point_itself=False)

df_distances = df.groupby(['Time']).apply(calculate_distances).reset_index()

и он выведет

    index  Time ID Group    X    Y  distance_Red  nearest_index_Red  \
0       0     1  A   Red  2.0  3.0      2.236068                  1   
1       1     1  B   Red  3.0  1.0      1.414214                  2   
2       2     1  C   Red  4.0  0.0      1.414214                  1   
3       3     1  X   Grn  2.0  0.0      1.414214                  1   
4       4     1  U   Grn  2.0  2.0      1.000000                  0   
5       5     1  V   Grn  1.0  1.0      2.000000                  1   
6       6     2  A   Red  1.0  2.0      3.162278                  8   
7       7     2  B   Red  6.0  0.0      2.236068                  8   
8       8     2  C   Red  4.0  1.0      2.236068                  7   
9       9     2  X   Grn  2.0  1.0      1.414214                  6   
10     10     2  U   Grn  5.0  0.0      1.000000                  7   
11     11     2  V   Grn  3.0  0.0      1.414214                  8   

    distance_Grn  nearest_index_Grn  
0       1.000000                  4  
1       1.414214                  4  
2       2.000000                  3  
3       1.414214                  5  
4       1.414214                  5  
5       1.414214                  3  
6       1.414214                  9  
7       1.000000                 10  
8       1.414214                 11  
9       1.414214                 11  
10      2.000000                 11  
11      1.414214                  9  

Я не пересчитывал и не проверял идентификаторы, но, по крайней мере, мне кажется, что он действительно возвращает идентификатор из подгруппы.