获取与数组元素相同的行的DataFrame索引

Question

我有一个DataFrame(temp(和一个数组(x(，它们的元素对应于DataFrame的一些行。我想获得DataFrame的索引，其对应的记录与数组的元素相同：

例如：

temp = pd.DataFrame({"A": [1,2,3,4], "B": [4,5,6,7], "C": [7,8,9,10]})
A   B   C
0   1   4   7
1   2   5   8
2   3   6   9
3   4   7   10
x = np.array([[1,4,7], [3,6,9]])

它应该返回索引：0和2。

我尝试了这个，但没有成功：

temp.loc[temp.isin(x[0])].index

Answer 1

使用numpy广播：

array = temp.to_numpy()[:, None]
mask = (array == x).all(axis=-1).any(axis=-1)
temp.index[mask]

Answer 2

我将转换为Multiindex，然后使用np.where转换为isin

i = pd.MultiIndex.from_frame(temp[['A','B','C']])
out = np.where(i.isin(pd.MultiIndex.from_arrays(x.T)))[0]

---

print(out)
#[0 2]

或者合并：

cols = ['A','B','C']
out = temp.reset_index().merge(pd.DataFrame(x,columns=cols)).loc[:,'index'].tolist()

或使用np.isin和all

out = temp.index[np.isin(temp[['A','B','C']],x).all(1)]

Answer 3

由于需要将DataFrame的整行与numpy array的行匹配，因此可以将DataFrame转换为数组，然后使用enumerate循环并返回索引：

temp_arr = temp.to_numpy()
for idx, row in enumerate(temp_arr):
if row in x:
print(idx)

输出：

0
2

使用列表理解的一种更优雅的方式是：

idx_list = [i for i, row in enumerate(temp_arr) if row in x ]    
print(idx_list)

输出：

[0, 2]