如何不使用for循环对DataFrame中的笛卡尔坐标数据进行Oversampling？

问题描述 #

我有一个包含X、Y、Z三列的3D pandas DataFrame，希望通过将每个(x,y)点的Z值替换为该点周围5单位范围内所有点的Z值均值，来进行Oversampling/平滑处理。我已通过如下for循环实现该功能：

import pandas as pd

Z_OS = []
X_OS = []
Y_OS = []
for inddex, row in df.iterrows():
    Z_OS += [df[(df['X'] > row['X']-5) & (df['X']<row['X']+5) & (df['Y'] > row['Y']-5) & (df1['Y']<row['Y']+5)]['Z'].mean()]
    X_OS += [row['X']]
    Y_OS += [row['Y']]

dict = {
    'X': X_OS,
    'Y': Y_OS,
    'Z': Z_OS
}
OSdf = pd.DataFrame.from_dict(dict)

但该方法处理大数据集时速度极慢，且不够Pythonic。请问如何不使用for loop实现该功能？是否可借助复杂的groupby函数完成？

解决方案 #

方法1：Scipy KDTree快速邻域查询（通用高效） #

这是处理空间邻域问题的最优方案之一，KDTree能快速定位每个点周围指定范围内的所有点，时间复杂度远低于循环遍历：

import pandas as pd
from scipy.spatial import KDTree

# 基于X、Y坐标构建KDTree
tree = KDTree(df[['X', 'Y']].values)

# 查询每个点周围5单位内的所有点索引
neighbor_indices = tree.query_ball_point(df[['X', 'Y']].values, r=5)

# 计算每个邻域内Z值的均值
df['Z_smoothed'] = [df.iloc[idx]['Z'].mean() for idx in neighbor_indices]

# 生成目标DataFrame
OSdf = df[['X', 'Y', 'Z_smoothed']].rename(columns={'Z_smoothed': 'Z'})

注意：如果数据集极大，可以考虑用numpy向量化操作进一步优化均值计算，避免列表推导的微小开销。

方法2：网格化分箱+GroupBy（适合规则分布数据） #

若你的X、Y坐标是规则网格分布的，可以通过分箱分组模拟邻域范围，利用groupby的高效性完成计算：

import pandas as pd
import numpy as np

# 按5单位间隔对X、Y分箱，确保覆盖所有点的邻域范围
x_bins = np.arange(df['X'].min() - 5, df['X'].max() + 5, 5)
y_bins = np.arange(df['Y'].min() - 5, df['Y'].max() + 5, 5)
df['X_bin'] = pd.cut(df['X'], bins=x_bins)
df['Y_bin'] = pd.cut(df['Y'], bins=y_bins)

# 计算每个分箱的Z均值
bin_z_means = df.groupby(['X_bin', 'Y_bin'])['Z'].mean().reset_index()

# 为每个点匹配其所在分箱及相邻分箱的均值（覆盖5单位范围）
def get_neighbor_group_mean(row):
    x_left = row['X_bin'].left
    y_left = row['Y_bin'].left
    # 筛选当前分箱上下左右共9个相邻分箱
    target_bins = bin_z_means[
        bin_z_means['X_bin'].apply(lambda b: b.left in [x_left-5, x_left, x_left+5]) &
        bin_z_means['Y_bin'].apply(lambda b: b.left in [y_left-5, y_left, y_left+5])
    ]
    return target_bins['Z'].mean()

df['Z_smoothed'] = df.apply(get_neighbor_group_mean, axis=1)
OSdf = df[['X', 'Y', 'Z_smoothed']].rename(columns={'Z_smoothed': 'Z'})

说明：该方法仅适用于规则分布的数据，对离散不规则的点集，KDTree的方案更可靠。

关于GroupBy的说明 #

单纯使用GroupBy无法直接实现“任意点周围5单位连续范围”的均值计算——GroupBy依赖离散分组，而你的需求是基于连续空间的邻域查询。但结合分箱操作（如方法2）可以间接模拟邻域效果，不过通用性和效率都不如KDTree方案。

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