编辑:为这个问题的所有解决方案添加时间
我更喜欢没有的解决方案apply
df['Flag'] = df.reset_index().melt(id_vars='index', value_name='val', var_name='col').query('val > 3').groupby('index')['col'].agg(list)
或者
df['Flag'] = df.stack().rename('val').reset_index(level=1).query('val > 3').groupby(level=0)['level_1'].agg(list)
Out[2576]:
A B C D E F Flag
0 1 4 7 1 5 7 [B, C, E, F]
1 2 5 8 3 3 4 [B, C, F]
2 3 6 9 5 6 3 [B, C, D, E]
测试数据:
a = [
[1, 4, 7, 1, 5, 7],
[2, 5, 8, 3, 3, 4],
[3, 6, 9, 5, 6, 3],
] * 10000
df = pd.DataFrame(a, columns = list('ABCDEF'))
计时%timeit
:
In [79]: %timeit (df>3).dot(df.columns).apply(list)
40.8 ms ± 1.66 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
In [80]: %timeit [df.columns[x].tolist() for x in df.gt(3).values]
1.23 s ± 10.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [81]: %timeit df.gt(3).dot(df.columns).apply(list)
37.6 ms ± 644 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
In [82]: %timeit df.T.apply(lambda x: list(x[x>3].index))
16.4 s ± 99.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [83]: %timeit df.stack().rename('val').reset_index(level=1).query('val > 3')
...: .groupby(level=0)['level_1'].agg(list)
4.05 s ± 15.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [84]: %timeit df.apply(lambda x: df.columns[np.argwhere(x>3).ravel()].values
...: , 1)
c:\program files\python37\lib\site-packages\numpy\core\fromnumeric.py:56: Future
Warning: Series.nonzero() is deprecated and will be removed in a future version.
Use Series.to_numpy().nonzero() instead
return getattr(obj, method)(*args, **kwds)
12 s ± 45.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
最快的是使用解决方案.dot