Pandas:按行从DataFrame的特定列中选择值

时间:2017-12-27 19:22:44

标签: python pandas numpy indexing

鉴于具有多列的DataFrame,我们如何逐行选择特定值来创建新系列?

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]

一种有效的方法是使用apply语句。

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

不幸的是,这不是矢量化操作,需要很长时间才能处理大型数据集。任何想法都将不胜感激。

2 个答案:

答案 0 :(得分:10)

Pandas approach

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

答案 1 :(得分:8)

NumPy方式

这是使用advanced indexing -

的矢量化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的列名进行排序,我们需要将sorter参数与np.searchsorted一起使用。为此类通用col_idx提取df的代码为:

# https://stackoverflow.com/a/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解决了一般情况,这可能是一个更好的选择,但是这篇文章中显示的其他可能的优化也可以很方便!