Pandas：计算 X、Y 分组之间的距离和角度

Question

Well, I have following columns:

Id PlayId  X     Y
0  0       2.3   3.4
1  0       5.4   3.2

2  1       3.2   5.1
3  1       4.2   1.7

If I have two rows groupped by one PlayId, I want to add two columns of Distance and Angle:

Id PlayId  X     Y   Distance_0  Distance_1 Angle_0 Angle_1
0  0       2.3   3.4 0.0         ?          0.0     ?
1  0       5.4   3.2 ?           0.0        ?       0.0

2  1       3.2   5.1
3  1       4.2   1.7

Every Distance-column describes Euclidean distance between i-th and j-th element in a group:

dist(x0, x1, y0, y1) = sqrt((x0 - x1) ** 2 + (y0 - y1) ** 2)

Similar way, the angle between i-th and j-th element is calculated.

So, how can I perform this efficiently, without processing elements one-by-one?

Answer 1

You can compute the pairwise distances by using the pdist function from SciPy:

df = pd.DataFrame({'X': [5, 6, 7], 'Y': [3, 4, 5]})

# df
#    X  Y
# 0  5  3
# 1  6  4
# 2  7  5

from scipy.spatial.distance import pdist, squareform

cols = [f'Distance_{i}' for i in range(len(df))]
pd.DataFrame(squareform(pdist(df.values)), columns=cols)

which produces the following DataFrame:

Distance_0  Distance_1  Distance_2
0   0.000000    1.638991    2.828427
1   1.638991    0.000000    1.638991
2   2.828427    1.638991    0.000000

This works, since pdist takes an array of size m * n, where m is the number of observations (=rows) and n the dimension of said observations (in this case: two - X and Y)

You could subsequently concat the original DataFrame with the newly created one if needed (using pd.concat ).

For the angle, you could use pdist as well, using metric='cosine' to compute the cosine distance. See this post for more information.