使用列数组矢量化 pandas 数据框列查找答案

【问题标题】：Vectorize pandas dataframe column lookup with array of columns使用列数组矢量化 pandas 数据框列查找
【发布时间】：2019-10-02 05:39:24
【问题描述】：

我有一个 Pandas 字符串数据框，如下所示。

import pandas as pd
import numpy as np
df = pd.DataFrame(np.random.randint(97,123,size=(3, 4), dtype=np.uint8).view('S1'), columns=list('ABCD'))
df

输出：

   A  B  C  D
0  q  g  v  f
1  l  m  u  u
2  r  r  j  w

我还有一个列名列表。

col_list = [['A'], ['A', 'B'], ['A', 'B', 'C']]

我想对df进行切片并应用如下操作：

df[col_list[1]].values.sum(axis=1)

输出：

array(['qg', 'lm', 'rr'], dtype=object)

同样，我需要对col_list 中的所有项目执行此操作。我可以在 for 循环中执行此操作，但如果列表很大，这会很慢。有什么方法可以对其进行矢量化处理，这样我就可以将col_list 作为一个numpy 数组传递，结果是一个形状为(len(col_list), len(df.index)) 的numpy 二维数组。

关键是，对于大型列表，它需要快速。

【问题讨论】：

标签： python pandas numpy vectorization

【解决方案1】：

将numpy 与r_、cumsum 和hsplit 一起使用：

import numpy as np

arr_list = np.hsplit(df.loc[:, np.r_[[i for l in col_list for i in l]]].values, 
               np.cumsum(list(map(len, col_list))))
res1 = list(map(lambda x:np.sum(x, 1), arr_list))[:-1]

如果col_list 有 3000 个列表，则比普通循环快 60 倍：

col_list = [['A'], ['A', 'B'], ['A', 'B', 'C']] * 1000

numpy:

%%timeit

arr_list = np.hsplit(df.loc[:, np.r_[[i for l in col_list for i in l]]].values, 
               np.cumsum(list(map(len, col_list))))
res1 = list(map(lambda x:np.sum(x, 1), arr_list))[:-1]
# 24.3 ms ± 3.36 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

for循环：

%%timeit

for l in col_list:
    df[l].values.sum(axis=1)
# 1.53 s ± 62.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

验证：

all(all(i == j) for i,j in zip(res1, res2))
# True

【讨论】：