使用idxmax函数的输出过滤pandas数据框

Question

我有一只熊猫。具有多个数值列的数据框架，并且希望找到跨行的最大值，因此我在下面做了

df = pd.DataFrame(np.random.random(size=(100000, 10)))
max_series = df.max(axis=1)
# O/P is a pd.Series like below
0     0.741459
1     0.995978
2     0.978618
3     0.973057
4     0.838006
...   

接下来我们要找到最大值的索引。所以我在下面写了

filter_ = df.idxmax(axis=1)
# O/P
0     3
1     8
2     7
3     5
4     1
..

现在在DataFrame上使用filter_，我想实现与max_series变量相同的结果，而不使用pd.DataFrame.max(axis=1)

所以我在下面试了试

df.loc[:, filter_]

或

df.filter(items=filter_, axis=1)

但是都给我

MemoryError: Unable to allocate 74.5 GiB for an array with shape (100000, 100000) and data type float64

我不需要100000x100000 matrix，我只需要max_series，也就是100000x1

那么我如何使用filter_过滤DataFrame并获得pd.Series的最大跨行?

Answer 1

这可能是一个更快的解决方案:

%%time
df = pd.DataFrame(np.random.random(size=(100000, 10)))
max_series = df.max(axis=1)
filter_  = df.idxmax(axis=1)
unique_cols = filter_.unique()
max_series_ = pd.concat([df.loc[df.index.isin(filter_[filter_ == col].index), col] for col in unique_cols]).sort_index()

from pandas.testing import assert_series_equal
assert_series_equal(max_series_, max_series)

也许还可以进一步优化。

Answer 2

这可能是解决方案之一，

filter_ = df.idxmax(axis=1)
df.apply(lambda row: row[filter_.loc[row.name]], axis=1)