在Pandas Series中使用布尔型Numpy数组进行向量化索引Numpy数组。

Question

以下可重现代码生成了一个示例数据集，模拟了我的数据，但规模小得多。

import numpy as np 
import pandas as pd

np.random.seed(142536)

df = pd.DataFrame({
        "vals": list(np.arange(12).reshape(3,4)),
        "idx" : list(np.random.choice([True, False], 12).reshape(3,4))})
df

                           idx            vals
0   [False, True, True, False]    [0, 1, 2, 3]
1    [True, True, False, True]    [4, 5, 6, 7] 
2  [False, True, False, False]  [8, 9, 10, 11] 

以下可重复代码返回我想要的结果，但对于大型数据集来说效率非常低。
如何更有效地完成此操作？

sel = []
for i in range(len(df.vals)):
    sel.append(df.vals[i][df.idx[i]])

df['sel'] = sel
df

                           idx            vals        sel
0   [False, True, True, False]    [0, 1, 2, 3]     [1, 2]
1    [True, True, False, True]    [4, 5, 6, 7]  [4, 5, 7]
2  [False, True, False, False]  [8, 9, 10, 11]        [9]

我尝试过np.apply_along_axis()、np.where()、df.apply()和df.transform()，但是在没有错误的情况下，无法让它们对这个案例起作用。

- Clay

---

vals中的所有行长度是否相同？ - user3483203

@user3483203 是的，它们是。 - Clay

Answer 1

前提条件不好，因为你不应该像这样存储数据。至少可以通过使用 itertools.chain 连接数据、索引数据，然后使用 np.array_split 将结果拆分来加快速度。

from itertools import chain

fn = lambda x: np.array(list(chain.from_iterable(x)))
df['sel'] = np.array_split(
    fn(df.vals)[fn(df.idx)], np.cumsum([sum(x) for x in df.idx][:-1]))

                           idx            vals      sel
0   [True, False, True, False]    [0, 1, 2, 3]   [0, 2]
1  [False, False, False, True]    [4, 5, 6, 7]      [7]
2   [False, True, True, False]  [8, 9, 10, 11]  [9, 10]

- cs95

Answer 2

使用列表推导和numpy索引：

df.assign(sel=[x[y] for x, y in zip(df.vals, df.idx)])

                           idx            vals      sel
0   [True, False, True, False]    [0, 1, 2, 3]   [0, 2]
1  [False, False, False, True]    [4, 5, 6, 7]      [7]
2   [False, True, True, False]  [8, 9, 10, 11]  [9, 10]

- user3483203

---

注意，df.vals 应该是一个数组而不是列表。如果它是一个列表，你需要使用 np.array(x)[y]。 - cs95

如果他的代码只是用于演示，并且实际上正在使用列表，那么这种方法就不会那么有用了。 - user3483203

我不认为这个推导式会比追加循环快得多。使用zip比范围索引更清洁一些。 - hpaulj

Answer 3

如果这是 df：

             vals                         idx
0    [0, 1, 2, 3]   [True, False, True, True]
1    [4, 5, 6, 7]  [False, True, False, True]
2  [8, 9, 10, 11]   [True, True, True, False]

那么你的 sel 就是：

In [21]: sel
Out[21]: [array([0, 2, 3]), array([5, 7]), array([ 8,  9, 10])]

这是一个大小不同的数组列表。

df列作为数组如下：

In [7]: vals = df['vals'].values
In [8]: idx = df['idx'].values

这两个都是对象数组的数组。但我们可以使用stack（或vstack）将它们转换为二维数组：

In [23]: vals = np.stack(vals)
In [24]: idx = np.stack(idx)
In [25]: vals
Out[25]: 
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
In [26]: idx
Out[26]: 
array([[ True, False,  True,  True],
       [False,  True, False,  True],
       [ True,  True,  True, False]])

我们可以使用布尔掩码进行索引，但结果是一个一维数组：

In [27]: vals[idx]
Out[27]: array([ 0,  2,  3,  5,  7,  8,  9, 10])

where 在 idx 上产生等效的索引数组元组：

In [28]: np.where(idx)
Out[28]: (array([0, 0, 0, 1, 1, 2, 2, 2]), array([0, 2, 3, 1, 3, 0, 1, 2]))

我们还可以从这些数组中生成一个掩码数组：

In [34]: mvals = np.ma.MaskedArray(vals, ~idx)
In [35]: mvals
Out[35]: 
masked_array(
  data=[[0, --, 2, 3],
        [--, 5, --, 7],
        [8, 9, 10, --]],
  mask=[[False,  True, False, False],
        [ True, False,  True, False],
        [False, False, False,  True]],
  fill_value=999999)
In [36]: mvals.compressed()
Out[36]: array([ 0,  2,  3,  5,  7,  8,  9, 10])

但是要逐行获取值，我们必须进行某种形式的迭代：

In [37]: [row[i] for row,i in zip(vals, idx)]
Out[37]: [array([0, 2, 3]), array([5, 7]), array([ 8,  9, 10])]

因此，In[7]和In[8]中的对象数组与堆叠的二维数组一样好，甚至更好。

In [40]: [row[i] for row,i in zip(df['vals'], df['idx'])]
Out[40]: [array([0, 2, 3]), array([5, 7]), array([ 8,  9, 10])]

而且你的range/append循环也很好（如果不是更好的话）。

你的sel数组大小变化（或者至少在理论上可以变化），这表明“向量化”的整个数组操作可能不可行。但是你需要这样一个列表吗？如果你不能用快速的数组操作生成它，那么你也不能用它进行操作。无论是创建还是使用，你都必须迭代“行”。

- hpaulj

---

选择这个作为答案是因为它有最详细的解释，并且基本上是我想出的相同的基础方法。不过，比较一下这里提出的所有解决方案的时间会很好。 - Clay

Answer 4

您不应该使用Pandas系列来存储列表。但是，如果避免不了，您可以使用itertools.compress和map，将df['vals']和df['idx']作为单独的参数传递：

from itertools import compress

df['sel'] = list(map(list, map(compress, df['vals'], df['idx'])))

print(df)

             vals                         idx        sel
0    [0, 1, 2, 3]   [False, True, True, True]  [1, 2, 3]
1    [4, 5, 6, 7]   [False, True, True, True]  [5, 6, 7]
2  [8, 9, 10, 11]  [True, False, False, True]    [8, 11]

如果您的df['vals']系列确实是NumPy数组，则可以使用NumPy索引：

df['sel'] = [vals[idx] for vals, idx in zip(df['vals'], df['idx'])]

- jpp

Answer 5

如果你将 Apply 函数解包成一个函数，它应该可以正常工作。至于任何速度上的提升，请在使用情况/数据方面向我们报告，因为反复调用该函数可能会非常昂贵：

def return_indices(row):
    row_vals = row['vals']
    row_idx = row['idx']
    true_rows = np.where(row_idx == True)
    return list(row_vals[true_rows])

df['sel'] = df.apply(lambda x: return_indices(x), axis=1)

- datawrestler

Answer 6

感谢大家的回答。

以下是我想出来的。我还没有与其他解决方案进行时间比较。

tmp = np.where(
        np.concatenate(df.idx.values).reshape(df.idx.values.shape[0],df.idx[0].shape[0] ), 
        np.concatenate(df.vals.values).reshape(df.vals.values.shape[0],df.vals[0].shape[0] ),
        np.nan)

df['sel'] = [*map(lambda a: [x for x in a if ~np.isnan(x)], tmp)]

df

                           idx            vals              sel
0   [False, True, True, False]    [0, 1, 2, 3]       [1.0, 2.0]
1    [True, True, False, True]    [4, 5, 6, 7]  [4.0, 5.0, 7.0]
2  [False, True, False, False]  [8, 9, 10, 11]            [9.0]

我认为这个lambda函数应该被映射(应用)到tmp np.array并且不需要追踪i的内部状态，所以它比我在原始帖子中提出的for循环更好（虽然我没有测试）。除非这就是python用for循环实现的方式。

编辑：

原始帖子中的for循环明显地更快。我没有确切的时间，但对于我的大型数据集，这个答案中的map函数需要几分钟才能完成，而OP中的for循环只需要几秒钟。

@hpaulj的评论：“而你的范围/附加循环几乎一样好（如果不是更好）”是正确的。

- Clay