在单个操作中使用二进制列表拆分 numpy 数组答案

【问题标题】：Split a numpy array with a binary list in a single operation在单个操作中使用二进制列表拆分 numpy 数组
【发布时间】：2019-07-08 08:16:55
【问题描述】：

我可以像这样将一个数组拆分成两个更小的数组：

>>> import numpy as np
>>> a = np.array([1,2,3,4,5])
>>> selector = np.array([True, False, True, True, False])
>>> selected, not_selected = a[selector], a[~ selector]
>>> selected
array([1, 3, 4])
>>> not_selected
array([2, 5])

但是，即使我在同一行生成 selected 和 not_selected，我 （至少我认为我是这样） 有效地在 a 上运行了两次，一次使用 @ 987654325@ 和它的倒数。如何使用真正单一且可能更快的 numpy 操作生成 selected 和 not_selected？或者这仍然是最好的方法吗？

【问题讨论】：

标签： python numpy split binary

【解决方案1】：

如果您对numba 持开放态度，我们可以获得一些内存效率并转移到显着的性能。提升 -

from numba import njit

@njit(parallel=True)
def select_numba(a, selector, out1, out2):
    iter1 = 0
    iter2 = 0
    for i,j in zip(a,selector):
        if j:
            out1[iter1] = i
            iter1 += 1
        else:
            out2[iter2] = i
            iter2 += 1
    return out1,out2

def select(a, selector):
    L = np.count_nonzero(selector)
    nL = len(selector)-L
    out1 = np.empty(L, dtype=a.dtype)
    out2 = np.empty(nL, dtype=a.dtype)
    select_numba(a,selector, out1, out2)        
    return out1,out2

示例运行 -

In [65]: a = np.array([1,2,3,4,5])
    ...: selector = np.array([True, False, True, True, False])

In [66]: select(a, selector)
Out[66]: (array([1, 3, 4]), array([2, 5]))

大型数据集的基准测试

In [60]: np.random.seed(0)
    ...: a = np.random.randint(0,9,(100000))
    ...: selector = np.random.rand(len(a))>0.5

In [62]: %timeit selected, not_selected = a[selector], a[~ selector]
1000 loops, best of 3: 1.2 ms per loop

In [63]: %timeit select(a, selector)
1000 loops, best of 3: 454 µs per loop

【讨论】：

我不知道 numba，所以对此表示赞同。但如果可能的话，我正在等待一个接受的答案，它只使用 numpy