Pythonic и эффективный способ сделать элемент "in" с помощью numpy
Я ищу способ эффективного получения массива булевых элементов, где заданы два массива с равным размером a
и b
, каждый элемент имеет значение true, если соответствующий элемент a
появляется in
соответствующий элемент b
.
Например, следующая программа:
a = numpy.array([1, 2, 3, 4])
b = numpy.array([[1, 2, 13], [2, 8, 9], [5, 6], [7]])
print(numpy.magic_function(a, b))
Должен печатать
[True, True, False, False]
Помните, что эта функция должна быть эквивалентна
[x in y for x, y in zip(a, b)]
Только numpy
-оптимизирован для случаев, когда a
и b
являются большими, и каждый элемент из b
достаточно мал.
Ответы
Ответ 1
Чтобы воспользоваться принципами NumPy broadcasting, сначала нужно создать квадрат квадрата b
, который может быть достигнут с помощью itertools.izip_longest
:
from itertools import izip_longest
c = np.array(list(izip_longest(*b))).astype(float)
в результате:
array([[ 1., 2., 5., 7.],
[ 2., 8., 6., nan],
[ 13., 9., nan, nan]])
Затем, выполнив np.isclose(c, a)
, вы получите 2D-массив логических переменных, показывающий разницу между каждыми c[:, i]
и a[i]
в соответствии с правилами вещания, предоставляя:
array([[ True, True, False, False],
[False, False, False, False],
[False, False, False, False]], dtype=bool)
Что можно использовать для получения ответа:
np.any(np.isclose(c, a), axis=0)
#array([ True, True, False, False], dtype=bool)
Ответ 2
Существует ли верхний предел длины небольших списков в b
? Если да, возможно, вы могли бы сделать b
матрицу, скажем 1000x5, и использовать nan
, чтобы заполнить пробелы для слишком малых субмассивов. Затем вы можете использовать numpy.any, чтобы получить ответ, который вы хотите:
In [42]: a = np.array([1, 2, 3, 4])
...: b = np.array([[1, 2, 13], [2, 8, 9], [5, 6], [7]])
In [43]: bb = np.full((len(b), max(len(i) for i in b)), np.nan)
In [44]: for irow, row in enumerate(b):
...: bb[irow, :len(row)] = row
In [45]: bb
Out[45]:
array([[ 1., 2., 13.],
[ 2., 8., 9.],
[ 5., 6., nan],
[ 7., nan, nan]])
In [46]: a[:,np.newaxis] == bb
Out[46]:
array([[ True, False, False],
[ True, False, False],
[False, False, False],
[False, False, False]], dtype=bool)
In [47]: np.any(a[:,np.newaxis] == bb, axis=1)
Out[47]: array([ True, True, False, False], dtype=bool)
Не знаю, быстрее ли это для ваших данных.
Ответ 3
Резюме
Подход Саулдо Кастро наиболее быстро проходит среди тех, кто был опубликован до сих пор. Выражение генератора в исходном столбце является вторым самым быстрым.
Код для генерации тестовых данных:
import numpy
import random
alength = 100
a = numpy.array([random.randint(1, 6) for i in range(alength)])
b = []
for i in range(alength):
length = random.randint(1, 5)
element = []
for i in range(length):
element.append(random.randint(1, 6))
b.append(element)
b = numpy.array(b)
print a, b
Параметры:
from itertools import izip_longest
def magic_function1(a, b): # From OP Martin Fixman
return [x in y for x, y in zip(a, b)]
def magic_function2(a, b): # What I thought might be better.
bools = []
for x, y in zip(a,b):
found = False
for j in y:
if x == j:
found=True
break
bools.append(found)
def magic_function3(a, b): # What I tried first
bools = []
for i in range(len(a)):
found = False
for j in range(len(b[i])):
if a[i] == b[i][j]:
found=True
break
bools.append(found)
def magic_function4(a, b): # From Bas Swinkels
bb = numpy.full((len(b), max(len(i) for i in b)), numpy.nan)
for irow, row in enumerate(b):
bb[irow, :len(row)] = row
a[:,numpy.newaxis] == bb
return numpy.any(a[:,numpy.newaxis] == bb, axis=1)
def magic_function5(a, b): # From Sauldo Castro, revised version
c = numpy.array(list(izip_longest(*b))).astype(float)
return numpy.isclose(c, a), axis=0)
Время n_executions
n_executions = 100
clock = timeit.Timer(stmt="magic_function1(a, b)", setup="from __main__ import magic_function1, a, b")
print clock.timeit(n_executions), "seconds"
# Repeat with each candidate function
Результаты:
- 0.158078225475 секунд для magic_function1
- 0.181080926835 секунд для magic_function2
- 0.259621047822 секунд для magic_function3
- 0.287054750224 секунд для magic_function4
- 0.0839162196207 секунд для magic_function5