python - pandas 系列中的矢量化索引 numpy 数组与 pandas 系列中的 bool 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()，但无法让它们中的任何一个在这种情况下正常工作。

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]