如何讓 numpy.where() 只返回滿足條件的元素？

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]
                  })

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

我試過了：

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

但這給出了以下錯誤：

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

我意識到我可以通過在where的else部分返回一些像 -1 這樣的默認值來解決這個問題，然后從ids中刪除所有出現的 -1 ，但這對於我巨大的 Dataframe 都無效並且似乎不是最優雅的方法。

如何以最有效（最省時）的方式解決這個問題？ 如果where不是最有效的解決方案，我願意接受其他建議。

Answer 1

您可以使用 boolean 索引或 dataframe 上的查詢方法在 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])