为每一行获取满足条件的列名列表的有效（可能是最短）方法答案

【问题标题】：Efficient (& perhaps shortest) way of getting a list of column names that satisfies a condition, for each row为每一行获取满足条件的列名列表的有效（可能是最短）方法
【发布时间】：2020-06-23 06:03:38
【问题描述】：

我有一个包含 5 列的数据框（第一列是 ID，4 列是国家/地区名称）

我想要每行的国家名称列表，满足相应国家列的特定条件。

df = {'id':['i1','i2','i3','i4','i5'], 'c1':[3,2,4,1,4], 'c2':[4,2,5,5,5], 'c3':[4,5,3,3,3], 'c4':[5,1,2,2,2]}

在上述情况下，我需要一个评级为 4 及以上的 ID。

我希望输出是每个 ID 的评级为 4 及以上的公司列表。可以是数据框或字典。

highest_rated_companies = { 'i1': ['c2', 'c3', 'c4'], 'i2': ['c3'], 'i3': ['c1', 'c2'], 'i4': ['c2'], 'i5': ['c1', 'c2'] }

【问题讨论】：

请分享数据而不是图片，在这种情况下是您的输入数据框
你试过什么？哪里出错了？输出是要被理解为您编写的字典还是再次隐式地作为数据框？
@对不起，我已经分享了数据
@DavidWierichs 输出可以是字典或数据框。我正在寻找一个简短的代码。我使用 2 个“for 循环”完成了它，但效率不高。

标签： python pandas numpy dictionary

【解决方案1】：

你可以尝试这样的事情，使用to_records，这似乎是你看到的最快的here：

第一个选项

import pandas as pd
import numpy as np
data = {'id':['i1','i2','i3','i4','i5'], 'c1':[3,2,4,1,4], 'c2':[4,2,5,5,5], 'c3':[4,5,3,3,3], 'c4':[5,1,2,2,2]}

df = pd.DataFrame(data)
print(df)
highest_rated_companies={row[1]:[df.columns[idx] for idx,val in enumerate(list(row)[2:],1) if val>=4] for row in df.to_records()}

第二个选项

import pandas as pd

data = {'id':['i1','i2','i3','i4','i5'], 'c1':[3,2,4,1,4], 'c2':[4,2,5,5,5], 'c3':[4,5,3,3,3], 'c4':[5,1,2,2,2]}

df = pd.DataFrame(data)
print(df)
highest_rated_companies={row[0]:[df.columns[idx] for idx,val in enumerate(row[1:],1) if val>=4] for i, row in df.iterrows()}
print(highest_rated_companies)

两个输出：

df:

   id  c1  c2  c3  c4
0  i1   3   4   4   5
1  i2   2   2   5   1
2  i3   4   5   3   2
3  i4   1   5   3   2
4  i5   4   5   3   2


highest_rated_companies:

{'i1': ['c2', 'c3', 'c4'], 'i2': ['c3'], 'i3': ['c1', 'c2'], 'i4': ['c2'], 'i5': ['c1', 'c2']}

时间戳：

第一个选项：

0.0113047 秒最佳情况，当执行时100 是脚本的倍数
1.2424291999999468 秒最佳情况，当执行时10000 是脚本的倍数

第二个选项

0.07292359999996734 秒最好的情况，当执行 100 时是脚本的倍数
7.821904700000005 秒最佳情况，当执行 10000 时脚本倍

编辑：

使用dt.to_records()，似乎是最快的方式，因为我测试了Ehsan的答案，当执行10000时我得到了脚本，时间戳为50.64001639999992秒，当执行100时我得到了脚本， 0.5399872999998934 秒的时间戳。即使它比 Second Option 更快，First Option 仍然是最快的。

【讨论】：

【解决方案2】：

你可以这样做：

df = pd.DataFrame(df)
keys, values = np.where(df[['c1','c2','c3','c4']].ge(4))
highest_rated_companies = pd.DataFrame({'id':df.iloc[keys].id,'c':df.columns[values+1]})

输出：

   id   c
0  i1  c2
0  i1  c3
0  i1  c4
1  i2  c3
2  i3  c1
2  i3  c2
3  i4  c2
4  i5  c1
4  i5  c2

如果您愿意，可以轻松地将其转换为字典。

【讨论】：

【解决方案3】：

另一种选择是使用to_dict 方法。如果您将 id 列设置为索引，您可以这样做：

df = df[df>=4]
d = df.to_dict('index')
output = {ID: [name for name,val in row.items() if not pd.isnull(val)] for ID, row in d.items()}

最后一行是将字典转换为所需的格式。时间测试：

In[0]:
import pandas as pd

df = {'id':['i1','i2','i3','i4','i5'], 'c1':[3,2,4,1,4], 'c2':[4,2,5,5,5], 'c3':[4,5,3,3,3], 'c4':[5,1,2,2,2]}

df = pd.DataFrame(df)
df = df.set_index('id',drop=True)
df = df[df>=4]

%%timeit -n 1000
d = df.to_dict('index')
output = {ID: [name for name,val in row.items() if not pd.isnull(val)] for ID, row in d.items()}

Out[0]
243 µs ± 23.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

虽然没有@MrNobody33 回答的那么快： 135 µs ± 4.29 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

【讨论】：