Question

我需要找到在Pandas DataFrame中获得最大值（每行）的所有索引。例如，如果我有这样的dataFrame：

   cat1  cat2  cat3
0     0     2     2
1     3     0     1
2     1     1     0

那么我正在寻找的方法会产生如下结果：

[['cat2', 'cat3'],
 ['cat1'],
 ['cat1', 'cat2']]

这是一个列表列表，但其他一些数据结构也没问题。

我无法使用df.idxmax(axis=1)，因为它只会产生第一个最大值。

Answer 1

以下是不同数据结构中的信息：

In [8]: df = pd.DataFrame({'cat1':[0,3,1], 'cat2':[2,0,1], 'cat3':[2,1,0]})

In [9]: df
Out[9]: 
   cat1  cat2  cat3
0     0     2     2
1     3     0     1
2     1     1     0

[3 rows x 3 columns]

In [10]: rowmax = df.max(axis=1)

最大值由真值表示：

In [82]: df.values == rowmax[:,None]
Out[82]: 
array([[False,  True,  True],
       [ True, False, False],
       [ True,  True, False]], dtype=bool)

np.where返回上面的DataFrame为True的索引。

In [84]: np.where(df.values == rowmax[:,None])
Out[84]: (array([0, 0, 1, 2, 2]), array([1, 2, 0, 0, 1]))

第一个数组表示axis=0的索引值，axis=1的第二个数组。每个数组中有5个值，因为有五个位置为True。

您可以使用itertools.groupby来构建您发布的列表列表，但考虑到上述数据结构，您可能不需要这样做：

In [46]: import itertools as IT

In [47]: import operator

In [48]: idx = np.where(df.values == rowmax[:,None])

In [49]: groups = IT.groupby(zip(*idx), key=operator.itemgetter(0))

In [50]: [[df.columns[j] for i, j in grp] for k, grp in groups]
Out[50]: [['cat1', 'cat1'], ['cat2'], ['cat3', 'cat3']]

Answer 2

你可以做到

In [2560]: cols = df.columns.values

In [2561]: vals = df.values

In [2562]: [cols[v].tolist() for v in vals == vals.max(1)[:, None]]
Out[2562]: [['cat2', 'cat3'], 
            ['cat1'], 
            ['cat1', 'cat2']]

<强>更新

以下是一个完整的例子：

import pandas as pd
import numpy as np

np.random.seed(400)

df = pd.DataFrame({
    'a': np.random.randint(0,3,size=10), 
    'b': np.random.randint(0,3,size=10), 
    'c': np.random.randint(0,5,size=10),
})

print(df)

out = [df.columns[i].tolist() for i in df.values == df.max(axis=1)[:,None]]

for i in out:
    print(i)

从print(df)返回：

来自print(out)：

['c']
['c']
['a', 'b', 'c']
['c']
['a', 'b']
['a', 'b', 'c']
['c']
['a', 'c']
['c']
['c']

在Pandas DataFrame中查找所有最大索引

2 个答案: