如何基于特征向量高效将pandas dataframe拆分为两个子集

报错原因 #

1. 你写的false_rows = g.get_group漏了参数调用，正确写法是g.get_group(False)
2. apply默认按列遍历（axis=0），你没加axis=1参数，实际处理的是列而不是行，得到的mask结果完全不符合预期
3. 逐行apply计算得到的布尔值是numpy.bool_类型，不是Python原生布尔值，groupby时会被识别为不可哈希的numpy数组类型，触发类型错误

最优实现方案（性能远高于逐行遍历版本） #

直接使用numpy批量向量化运算，完全避免逐行操作，比你最初的iterrows版本快10~100倍：

import numpy as np

def partition(dataset, question):
    # 将特征列堆叠为二维数组，形状为 [样本数, 特征维度]
    feature_matrix = np.stack(dataset.iloc[:, 0].values)
    # 批量计算所有样本和question向量的点积，生成布尔掩码
    mask = (feature_matrix @ question.vector) >= 1
    # 直接用布尔索引拆分数据集
    true_rows = dataset[mask].copy()
    false_rows = dataset[~mask].copy()
    return true_rows, false_rows

该方案没有多余的中间列，也不会修改原数据集，性能最高。

保留groupby思路的修正版 #

def partition(dataset, question):
    df = dataset.copy()
    # 新增axis=1按行遍历，同时把numpy布尔值转成Python原生bool解决哈希报错
    mask = df.apply(lambda x: bool(np.dot(x[0], question.vector) >= 1), axis=1)
    df['mask'] = mask
    g = df.groupby('mask')
    # 提取分组后删掉多余的mask列
    true_rows = g.get_group(True).drop('mask', axis=1)
    false_rows = g.get_group(False).drop('mask', axis=1)
    return true_rows, false_rows

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