Python Pandas DataFrame每行内各元素到最近邻的距离计算方法

思路说明 #

要计算每行内每个元素的最近邻距离，本质是对每个位置的元素，取同一行内除自身外所有元素与该元素的绝对差的最小值。

代码实现 #

方法1：逐行处理（可读性高，适合小数据量） #

import pandas as pd

# 构造示例DataFrame
d = {'col1': [1, 2], 'col2': [3, 4], 'col3': [4, 10], 'col4': [1, 8]}
df = pd.DataFrame(data=d)

# 定义单行最近邻距离计算函数
def calc_row_nn(row):
    res = []
    for idx, val in enumerate(row):
        # 排除当前元素本身，计算和其他所有元素的绝对差
        other_vals = row.drop(row.index[idx])
        min_diff = abs(other_vals - val).min()
        res.append(min_diff)
    return pd.Series(res, index=row.index)

# 按行应用函数得到结果
df_nn_dist = df.apply(calc_row_nn, axis=1)
print(df_nn_dist)

输出结果：

col1  col2  col3  col4
0     0     1     1     0
1     2     2     2     2

方法2：向量化计算（性能高，适合大数据量） #

如果处理的DataFrame规模很大，可以用numpy广播实现向量化运算，避免循环带来的性能损耗：

import numpy as np
import pandas as pd

d = {'col1': [1, 2], 'col2': [3, 4], 'col3': [4, 10], 'col4': [1, 8]}
df = pd.DataFrame(data=d)
arr = df.values

# 广播计算每行内所有元素两两绝对差
diff_mat = np.abs(arr[:, :, None] - arr[:, None, :])
# 把元素和自身的差值（对角线位置）替换为无穷大，避免干扰最小值计算
np.fill_diagonal(diff_mat, np.inf, axis1=1, axis2=2)
# 取每个元素对应的最小差值
nn_dist_arr = diff_mat.min(axis=2)

df_nn_dist = pd.DataFrame(nn_dist_arr, index=df.index, columns=df.columns)
print(df_nn_dist)

输出结果和方法1完全一致。

内容的提问来源于stack exchange，提问作者MathMan 99