Question

我想使用 numpy.where() 遍历 Pandas DataFrame 并获取一个仅包含满足条件的元素的列表。

例如，假设我有以下 Pandas DataFrame：

df = pd.DataFrame({"A": [1, 2, 3, 5, 3, 7, 3],
                   "B": [0, 1, 6, 4, 9, 8, 2],
                   "id": [0, 1, 2, 3, 4, 5, 6]
                  })

我想返回列 id 等于 3 且列 A 大于或等于 5 的那些 B 值的列表

我试过了：

ids = np.where((df["A"] == 3) & (df["B"] >= 5)), df["id"])

但是会出现以下错误：

ValueError: either both or neither of x and y should be given

我意识到我可以通过在 else 的 where 部分返回一些默认值（如 -1）然后从 ids 中删除所有出现的 -1 来解决这个问题，但那是对我庞大的 Dataframe 来说既无效又似乎不是最优雅的方式。

如何以最有效（耗时最少）的方式解决这个问题？如果 where 不是最有效的解决方案，我愿意接受其他建议。

Answer 1

您可以在 Pandas 中使用布尔索引或数据帧上的查询方法来完成此操作。

In [4]: import pandas as pd

In [5]: df = pd.DataFrame({"A": [1, 2, 3, 5, 3, 7, 3],
   ...:                    "B": [0, 1, 6, 4, 9, 8, 2],
   ...:                    "id": [0, 1, 2, 3, 4, 5, 6]
   ...:                   })

In [6]: df
Out[6]:
   A  B  id
0  1  0   0
1  2  1   1
2  3  6   2
3  5  4   3
4  3  9   4
5  7  8   5
6  3  2   6

In [7]: df[(df["A"] == 3) & (df["B"] >= 5)]['id'].to_list()
Out[7]: [2, 4]

In [8]: df.query("A == 3 and B >= 5")['id'].to_list()
Out[8]: [2, 4]

Answer 2

使用：

In [1225]: df.loc[(df["A"] == 3) & (df["B"] >= 5), 'id'].to_numpy()
Out[1225]: array([2, 4])

如何让 numpy.where() 只返回满足条件的元素？

2 个答案: