时间戳

首页 > 解决方案 > pandas 系列中的矢量化索引 numpy 数组与 pandas 系列中的布尔 numpy 数组

python - pandas 系列中的矢量化索引 numpy 数组与 pandas 系列中的布尔 numpy 数组

问题描述

以下可重现的代码生成了一个示例数据集，该数据集在更小范围内模拟了我的数据。

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()，但无法让它们中的任何一个在没有错误的情况下为这种情况工作。

标签： pythonpandasnumpy

解决方案

前提很糟糕，因为您不应该像这样存储数据。您至少可以通过使用 连接数据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]

推荐阅读