基于列值的高效逐行操作（聚合）

Question

我正在处理一个大型数据集，需要将某些列组合成一个列表。挑战在于，选择的列数取决于另一个键值，每行都不同。

示例：

原始数据集：

Key Col1 Col2 Col3 Col4 Col5 NonrelatedCols
1   a    b    c    d    e    ...
3   r    b    x    d    c    ...
2   k    d    q    l    w    ...
0   w    a    c    s    w    ...

预期结果：

Key  List      NonrelatedCols
1    [a]       ...
3    [r,b,x]   ...
2    [k,d]     ...
0    []        ...

主要挑战在于代码的性能。遍历行，不是最佳方法，我正在寻找更有效的方法。

我尝试只选择目标列，然后使用agg(list,axis=1) 函数。问题是，没有办法根据键值聚合列。看来，我仍然需要对所有行进行迭代，然后在合并列中删除过多的值

Cols = ['Col1', 'Col2', 'Col3','Col4','Col5']
CombinedList = rawData[Cols].agg(list,axis=1)
for i in range(rawData.shape[0]):
    CombinedList[i]=CombinedList[i][:rawData['Key'][i]]

总结一下，有没有办法简化代码，所以我不需要遍历所有的行？

Answer 1

只考虑重要的列：

    Key Col1 Col2 Col3 Col4 Col5
0    1    a    b    c    d    e
1    3    r    b    x    d    c
2    2    k    d    q    l    w
3    0    w    a    c    s    w

假设Key总是在第一列，applylist函数对除第一列（Key）之外的所有Col感兴趣的列，然后根据@987654327中的值索引列表@：

df.apply(lambda x : list(x[1:])[:x[0]], axis=1)

• 这要求Key 是int
  • 如果Key 不是int，则使用df = df.astype({'Key': 'int32'})

输出：

0          [a]
1    [r, b, x]
2       [k, d]
3           []

Answer 2

试试：

df.apply(lambda x: x[1:x.Key+1].tolist(), axis=1)

输出：

0          [a]
1    [r, b, x]
2       [k, d]
3           []
dtype: object