如何从Pandas DataFrame的每一行中选择特定列？

Question

我有一个以这种格式呈现的DataFrame：

    a   b   c
0   1   2   3
1   4   5   6
2   7   8   9
3   10  11  12
4   13  14  15

并且有一个像这样的数组，带有列名：

['a', 'a', 'b', 'c', 'b']

我希望从每一行中提取一个数据数组，每个数组元素对应一列。给定的列名数组决定了每行中需要提取哪一列的数据。在这个例子中，返回的结果为：

[1, 4, 8, 12, 14]

使用Pandas是否可以作为单个命令完成此操作，还是需要迭代？我尝试过使用索引。

i = pd.Index(['a', 'a', 'b', 'c', 'b'])
i.choose(df)

但我遇到了一个段错误，由于文档不足，我无法诊断问题。

- gggritso

Answer 1

您可以使用lookup函数。

>>> i = pd.Series(['a', 'a', 'b', 'c', 'b'])
>>> df.lookup(i.index, i.values)
array([ 1,  4,  8, 12, 14])

如果需要，i.index 可以与 range(len(i)) 不同。

- DSM

---

太棒了，谢谢！也可以对这些索引进行赋值吗？ - gggritso

你可以进行赋值，但仅当框架是单个dtype（就像现在这样）时。df.unstack().loc[zip(i.values,i.index)] = [1,2,3,4,5]。并且你必须在两侧匹配长度（你也可以使用此语法进行选择）；请参见此问题：https://github.com/pydata/pandas/issues/7138 - Jeff

如果您想要添加索引，请创建一个序列：pd.Series(df.lookup(i.index, i.values), index=i.index)。 - user394430

在 pandas 1.2.0 版本中，lookup 函数已被弃用，建议使用 .loc 或 .melt（请参见：https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.lookup.html）。 - MorningGlory

Answer 2

对于大型数据集，你可以在numpy的基本数据上使用索引，只需将列名转换为数值索引（在这种情况下很简单）：

df.values[arange(5),[0,0,1,2,1]]

out: array([ 1,  4,  8, 12, 14])

这将比列表推导或其他显式迭代更加高效。

- mdurant

---

这应该是新的被接受的答案。因为pd.lookup()现在已经被弃用，而melt()解决方案可能会在处理大型数据集时导致内存问题。 - Arturo Rodriguez

Answer 3

正如MorningGlory在评论中所述，lookup已在版本1.2.0中被弃用。

文档说明可以使用melt和loc实现相同的功能，但我认为这并不是很明显，因此在这里介绍一下。

首先，使用melt创建一个查找DataFrame：

i = pd.Series(["a", "a", "b", "c", "b"], name="col")
melted = pd.melt(
    pd.concat([i, df], axis=1),
    id_vars="col",
    value_vars=df.columns,
    ignore_index=False,
)

  col variable  value
0   a        a      1
1   a        a      4
2   b        a      7
3   c        a     10
4   b        a     13
0   a        b      2
1   a        b      5
2   b        b      8
3   c        b     11
4   b        b     14
0   a        c      3
1   a        c      6
2   b        c      9
3   c        c     12
4   b        c     15

然后，使用loc仅获取相关值：

result = melted.loc[melted["col"] == melted["variable"], "value"]

0     1
1     4
2     8
4    14
3    12
Name: value, dtype: int64

最后，如果需要，可以按照之前的索引顺序获取相同的索引顺序：

result.loc[df.index]

0     1
1     4
2     8
3    12
4    14
Name: value, dtype: int64

Pandas在文档中使用`factorize`和`numpy`索引提供了不同的解决方案： 文档链接。

df = pd.concat([i, df], axis=1)
idx, cols = pd.factorize(df['col'])
df.reindex(cols, axis=1).to_numpy()[np.arange(len(df)), idx]

[ 1  4  8 12 14]

- spettekaka

Answer 4

你可以随时使用列表推导式：

[df.loc[idx, col] for idx, col in enumerate(['a', 'a', 'b', 'c', 'b'])]

- Gregor

---

这不是矢量化的，你可以使用for循环来进行任何操作。 - Wildhammer

这个方法可以行得通，但与访问NumPy数组相比会非常慢，就像@mdurant的答案所述：https://dev59.com/5GAf5IYBdhLWcg3wSQ4z#24833522 - Reslan Tinawi