python - 如何根据不同的输入自动迭代以获取(pandas)数据帧的元素

Question

我有一个包含大数据集的项目，我正在尝试提取一些值，这些值取决于用户给出的输入，这意味着它彼此不同。我的数据框有 16 列，包括物种、重量、人口、颜色、位置等，看起来像(我为这个问题做了一个小例子)

df = pd.DataFrame({'species': ['bear', 'bear', 'marsupial', 'polar bear', 'bear', 'polar bear'],
                  'weights': [350, 350, 140, 450, 350, 540],
                  'population': [1864, 22000, 80000, 3000, 7864, 5000],
                  'color': ['brown', 'black', 'gray', 'white', 'brown', 'brown'],
                  'locations': ['US', 'Canada', 'Canada', 'Russia', 'US', 'Canada']})
output: 
   color locations  population species  weights
0  brown        US        1864    bear      350
1  black    Canada       22000    bear      350
4  brown        US        7864    bear      350

我询问用户您想查看哪些功能，然后从数据框中返回名称。我将用户输入收集为字典:

dict = {
  species: bear,
  weights: 350
}

我可以使用 loc 手动获取名称

df_loc = df.loc[(df['weights'] == 350) & (df['species'] == 'bear')]

output:
color locations  population species  weights
0  brown        US        1864    bear      350
1  black    Canada       22000    bear      350
4  brown        US        7864    bear      350

到目前为止，还好，但我无法找到一种方法将 df.loc[] 放入字典循环中，并且可以自动执行搜索过程，因为其中的大小和键值字典随时都会改变。

有什么想法吗？

Answer 1

您可以定义一个函数，它接受用户提供的特征字典并返回所需的数据帧。

代码:

def get_df(features):
    mask = True
    for k, v in features.items():
        mask = mask & (df[k] == v)
    return df.loc[mask]

输入功能:

features = {
  "species": "bear",
  "weights": 350
}

输出:

>>> print(get_df(features))

  species  weights  population  color locations
0    bear      350        1864  brown        US
1    bear      350       22000  black    Canada
4    bear      350        7864  brown        US