从列值与数组中的值匹配的数据框中提取具有重复的行

Question

我有一个 pandas.Dataframe df，其中 headers 列之一是 'X'。假设这是大小 (N,M)。 N=3,M=2 在这个例子中：

  X Y
0 1 a
1 2 b
2 3 c

我有一个大小为 (Q,) 的一维 numpy.array arr，它包含值，其中一些是重复的。本例中 Q=5：

array([1, 2, 3, 2, 2])

我想创建一个新的 pandas.Dataframe df_op，其中包含来自 df 的行，其中每个 row.X 匹配来自 arr 的条目。这意味着某些行被多次提取，结果 df_op 的大小为 (Q,M)。如果可能的话，我也想保持与 arr 中相同的条目顺序。

使用通常的布尔索引是行不通的，因为它只会选取唯一的行。我也想尽可能避免循环，因为 Q 很大。

如何获得df_op？谢谢。

Answer 1

使用索引多次获取同一行：

x = [1, 2, 3, 2, 2]
df = pd.DataFrame({'X': [1, 2, 3], 'Y': ['a', 'b', 'c']})

out = df.set_index('X').loc[x].reset_index()

输出：

>>> out
   X  Y
0  1  a
1  2  b
2  3  c
3  2  b
4  2  b

Extract rows with repeats, from dataframe where column value matches value from an array