Разделить массив с двоичным списком в одной операции - PullRequest
Новая
       23

Разделить массив с двоичным списком в одной операции

4 голосов
/ 08 июля 2019

Я могу разбить массив на два меньших массива, как это: 

>>> 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 дважды, один раз с selector и снова с обратным. Как бы я сгенерировал selected и not_selected, используя по-настоящему одиночную и, предположительно, более быструю, пустую операцию? Или это все еще лучший способ сделать это?

1 Ответ

2 голосов
/ 08 июля 2019

Если вы открыты для 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
...