python - 组名称在多列中时， Pandas 的前n个值

尝试从较大的DataFrame中查找前n个值。键是前两列中名称相似的对象的组合。但是，无论键位于哪一列，我都希望找到最大值。通过示例更好地演示：

import itertools
import pandas as pd

np.random.seed(10)

pairs = [combo for combo in itertools.combinations(['apple','banana','pear','orange'], 2)]

df = pd.DataFrame(pairs, columns=['a','b'])
df['score'] = np.random.rand(6)

原始的DataFrame：

In [2]: df
Out[2]: a   b   score
     0  apple   banana  0.771321
     1  apple   pear    0.020752
     2  apple   orange  0.633648
     3  banana  pear    0.748804
     4  banana  orange  0.498507
     5  pear    orange  0.224797

假设我有一个名为fruits的数据库表，可以模拟上面的df，这就是使用SQL完成我的任务的方式：

uniq = pd.unique(df[['a', 'b']].values.ravel())

df_sql = pd.DataFrame()
for fruit in uniq:
    dfsql_tmp = pd.read_sql_query(
    """SELECT a,b,score FROM fruits
    WHERE a = %s
    OR b = %s
    ORDER BY score DESC
    LIMIT 1;""",
    engine, params=[fruit, fruit])

    df_sql = pd.concat([df_sql, dfsql_tmp], ignore_index=True)

这正好满足了我的要求，每个唯一值（来自df['a']和df['b']的并集）的前n个得分。所需的输出：

In [5]: df_sql
Out[5]: a   b   score
     0  apple   banana  0.771321 #highest apple score
     1  apple   banana  0.771321 #highest banana score
     2  apple   orange  0.633648 #highest orange score
     3  banana  pear    0.748804 #highest pear score

编辑

这也可以解决问题，但是规模缓慢：

N=1
df_new = pd.DataFrame()
for fruit in uniq:
    df_tmp = df[(df['a'] == fruit) | (df['b'] == fruit)].sort_values('score', ascending=False).head(N)
    df_new= pd.concat([df_new, df_tmp])

有没有更好的方法来获得我想要的结果？嵌套的sql查询无法缩放。我宁愿在一个大df上执行操作。保持n而不是简单地设置max或min也很重要。

最佳答案

这不是一个很好的解决方案，我怀疑那里有更好的解决方案，但这是一个难题。这将创建一个约55万行x 5列的DataFrame，并在我的笔记本电脑上运行约4秒钟。

import string
import pandas as pd
import numpy as np
import itertools

np.random.seed(10)
pairs = [combo for combo in itertools.combinations(string.letters + string.digits, 4)]

df = pd.DataFrame(pairs, columns=['a', 'b', 'c', 'd'])
df['score'] = np.random.rand(len(df))

cols = ['a', 'b', 'c', 'd']
indexes = []

for c in pd.concat([df[col] for col in cols]).unique():
    indexes.append(df[reduce(lambda x, y: x | y, [df[col] == c for col in cols])]['score'].idxmax())
print df.ix[indexes]

如果您不希望原始索引保留在输出中，请在末尾添加.reset_index()。

对于前N个，而不是执行.idxmax()，请对缩小的帧进行排序，并使用.iloc[:N]获取前N个索引。