Pandas：按行选择DataFrame特定列的值

Question

在一个包含多列的 DataFrame 中，我们如何按行选择特定列的值来创建新的 Series？

df = pd.DataFrame({"A":[1,2,3,4], 
                   "B":[10,20,30,40], 
                   "C":[100,200,300,400]})
columns_to_select = ["B", "A", "A", "C"]

目标：[10, 2, 3, 400]

其中一种可行的方法是使用应用语句。

df["cols"] = columns_to_select
df.apply(lambda x: x[x.cols], axis=1)

很遗憾，这不是一个向量化的操作，在大型数据集上需要花费很长时间。欢迎提出任何想法。

- Jason Sanchez

Answer 1

Pandas 的方法:

In [22]: df['new'] = df.lookup(df.index, columns_to_select)

In [23]: df
Out[23]:
   A   B    C  new
0  1  10  100   10
1  2  20  200    2
2  3  30  300    3
3  4  40  400  400

- MaxU - stand with Ukraine

---

一秒钟落后你。;-) - BENY

@Wen，是的，我知道这种感觉 - 抱歉 :) - MaxU - stand with Ukraine

@MaxU 这正是我在寻找的。谢谢！ - Jason Sanchez

@JasonSanchez，很高兴我能帮到你 :) - MaxU - stand with Ukraine

几年后，pd.lookup已被弃用，以下是推荐的方法： https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#looking-up-values-by-index-column-labels - Tomer Cagan

Answer 2

NumPy方式

这里是使用高级索引的向量化NumPy方式 -

# Extract array data
In [10]: a = df.values

# Get integer based column IDs
In [11]: col_idx = np.searchsorted(df.columns, columns_to_select)

# Use NumPy's advanced indexing to extract relevant elem per row
In [12]: a[np.arange(len(col_idx)), col_idx]
Out[12]: array([ 10,   2,   3, 400])

如果 df 的列名没有排序，我们需要使用 np.searchsorted 中的 sorter 参数。提取这样一个通用的 df 的 col_idx 的代码如下：

# https://dev59.com/sWcs5IYBdhLWcg3wPxce#38489403/ @Divakar
def column_index(df, query_cols):
    cols = df.columns.values
    sidx = np.argsort(cols)
    return sidx[np.searchsorted(cols,query_cols,sorter=sidx)]

所以，col_idx 可以这样获得 -

col_idx = column_index(df, columns_to_select)

进一步优化

对其进行优化后发现，处理np.searchsorted字符串时成为瓶颈，通常情况下 NumPy 在处理字符串方面并不擅长。所以，为了解决这个问题并利用列名为单个字母的特殊情况，我们可以将它们快速转换为数字，然后将数字输入searchsorted以实现更快速的处理。

因此，在列名为单个字母且已排序的情况下获取基于整数的列 ID 的优化版本如下 -

def column_index_singlechar_sorted(df, query_cols):
    c0 = np.fromstring(''.join(df.columns), dtype=np.uint8)
    c1 = np.fromstring(''.join(query_cols), dtype=np.uint8)
    return np.searchsorted(c0, c1)

这给我们提供了一个修改过的解决方案，如下所示 -

a = df.values
col_idx = column_index_singlechar_sorted(df, columns_to_select)
out = pd.Series(a[np.arange(len(col_idx)), col_idx])

时间 -

In [149]: # Setup df with 26 uppercase column letters and many rows
     ...: import string
     ...: df = pd.DataFrame(np.random.randint(0,9,(1000000,26)))
     ...: s = list(string.uppercase[:df.shape[1]])
     ...: df.columns = s
     ...: idx = np.random.randint(0,df.shape[1],len(df))
     ...: columns_to_select = np.take(s, idx).tolist()

# With df.lookup from @MaxU's soln
In [150]: %timeit pd.Series(df.lookup(df.index, columns_to_select))
10 loops, best of 3: 76.7 ms per loop

# With proposed one from this soln
In [151]: %%timeit
     ...: a = df.values
     ...: col_idx = column_index_singlechar_sorted(df, columns_to_select)
     ...: out = pd.Series(a[np.arange(len(col_idx)), col_idx])
10 loops, best of 3: 59 ms per loop

鉴于 df.lookup 解决了一个通用的情况，这可能是更好的选择，但本文展示的其他可能的优化也很方便！

- Divakar