以下是不同数据结构中的信息:
In [8]: df = pd.DataFrame({'cat1':[0,3,1],'cat2':[2,'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)
最大值由True值表示:
In [82]: df.values == rowmax[:,None]
Out[82]:
array([[False,True,True],[ True,False,False],False]],dtype=bool)
np.where返回上面的DataFrame为True的索引.
In [84]: np.where(df.values == rowmax[:,None])
Out[84]: (array([0,2,2]),array([1,1]))
第一个数组表示轴= 0的索引值,第二个数组表示轴= 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']]