Эквивалент Python функции MATLAB "ismember"

После многих попыток оптимизации кода, кажется, что одним из последних ресурсов будет попытка выполнить код ниже с использованием нескольких ядер. Я не знаю точно, как преобразовать/переструктурировать мой код, чтобы он мог работать намного быстрее, используя несколько ядер. Буду признателен, если я смогу получить руководство для достижения конечной цели. Конечной целью является возможность быстрого запуска этого кода для массивов A и B, где каждый массив содержит около 700 000 элементов. Вот код с использованием небольших массивов. Массивы элементов 700k закомментированы.

import numpy as np

def ismember(a,b):
    for i in a:
        index = np.where(b==i)[0]
        if index.size == 0:
            yield 0
        else:
            yield index


def f(A, gen_obj):
    my_array = np.arange(len(A))
    for i in my_array:
        my_array[i] = gen_obj.next()
    return my_array


#A = np.arange(700000)
#B = np.arange(700000)
A = np.array([3,4,4,3,6])
B = np.array([2,5,2,6,3])

gen_obj = ismember(A,B)

f(A, gen_obj)

print 'done'
# if we print f(A, gen_obj) the output will be: [4 0 0 4 3]
# notice that the output array needs to be kept the same size as array A.

То, что я пытаюсь сделать, это имитировать функцию MATLAB под названием ismember [2] (тот, который отформатирован как: [Lia,Locb] = ismember(A,B) Я просто пытаюсь получить только часть Locb.

Из Matlab: Locb, содержащий наименьший индекс в B для каждого значения в A, являющегося членом B. Выходной массив, Locb, содержит 0, где A не является членом B

Одна из основных проблем заключается в том, что мне нужно максимально эффективно выполнять эту операцию. Для тестирования у меня есть два массива из 700 тыс. Элементов. Создание генератора и прохождение значений генератора, похоже, не ускоряет работу.

ОТВЕТЫ

Ответ 1

Прежде чем беспокоиться о нескольких ядрах, я бы устранил линейное сканирование в вашей функции ismember с помощью словаря:

def ismember(a, b):
    bind = {}
    for i, elt in enumerate(b):
        if elt not in bind:
            bind[elt] = i
    return [bind.get(itm, None) for itm in a]  # None can be replaced by any other "not in b" value

В исходной реализации требуется полное сканирование элементов в B для каждого элемента в A, что делает его O(len(A)*len(B)). Вышеприведенный код требует одного полного сканирования B для генерации dict Bset. Используя dict, вы эффективно выполняете поиск каждого элемента в константе B для каждого элемента A, делая операцию O(len(A)+len(B)). Если это все еще слишком медленно, то беспокойтесь о том, чтобы выполнить вышеуказанную функцию на нескольких ядрах.

Изменить: я также немного изменил вашу индексацию. Matlab использует 0, потому что все его массивы начинаются с индекса 1. Начальные массивы Python/numpy равны 0, поэтому, если ваш набор данных выглядит как

A = [2378, 2378, 2378, 2378]
B = [2378, 2379]

и вы возвращаете 0 без элемента, тогда ваши результаты будут исключать все элементы A. Вышеуказанная процедура возвращает None без индекса, а не 0. Возврат -1 является опцией, но Python будет интерпретировать это как последний элемент в массиве. None вызовет исключение, если оно используется как индекс в массиве. Если вы хотите различное поведение, измените второй аргумент в выражении Bind.get(item,None) на значение, которое вы хотите вернуть.

Ответ 2

sfstewman отличный ответ, скорее всего, решил проблему для вас.

Я просто хотел бы добавить, как вы можете достичь того же исключительно в numpy.

Я использую numpy unique функции in1d.

B_unique_sorted, B_idx = np.unique(B, return_index=True)
B_in_A_bool = np.in1d(B_unique_sorted, A, assume_unique=True)
  • B_unique_sorted содержит уникальные значения в B, отсортированные.
  • B_idx для этих значений содержит индексы в оригинале B.
  • B_in_A_bool - это булев массив размером B_unique_sorted, который сохраняет ли значение в B_unique_sorted в A.
    Примечание: Мне нужно искать (уникальные vals из B) в A, потому что мне нужен вывод, который нужно вернуть по отношению к B_idx
    Примечание: Я предполагаю, что A уже уникален.

Теперь вы можете использовать B_in_A_bool для получения общих vals

B_unique_sorted[B_in_A_bool]

и их соответствующие индексы в оригинале B

B_idx[B_in_A_bool]

Наконец, я предполагаю, что это значительно быстрее, чем чистый Python for-loop, хотя я его не тестировал.

Ответ 3

Попробуйте использовать понимание списка;

In [1]: import numpy as np

In [2]: A = np.array([3,4,4,3,6])

In [3]: B = np.array([2,5,2,6,3])

In [4]: [x for x in A if not x in B]
Out[4]: [4, 4]

Как правило, перечислимые функции намного быстрее, чем для циклов.

Чтобы получить равный список длин,

In [19]: map(lambda x: x if x not in B else False, A)
Out[19]: [False, 4, 4, False, False]

Это довольно быстро для небольших наборов данных:

In [20]: C = np.arange(10000)

In [21]: D = np.arange(15000, 25000)

In [22]: %timeit map(lambda x: x if x not in D else False, C)
1 loops, best of 3: 756 ms per loop

Для больших наборов данных вы можете попробовать использовать multiprocessing.Pool.map() для ускорения операции.

Ответ 4

Вот точный эквивалент MATLAB, который возвращает оба выходных аргумента [Lia, Locb], которые соответствуют MATLAB, за исключением Python 0, также является допустимым индексом. Таким образом, эта функция не возвращает 0s. Он по существу возвращает Locb (Locb > 0). Производительность также эквивалентна MATLAB.

def ismember(a_vec, b_vec):
    """ MATLAB equivalent ismember function """

    bool_ind = np.isin(a_vec,b_vec)
    common = a[bool_ind]
    common_unique, common_inv  = np.unique(common, return_inverse=True)     # common = common_unique[common_inv]
    b_unique, b_ind = np.unique(b_vec, return_index=True)  # b_unique = b_vec[b_ind]
    common_ind = b_ind[np.isin(b_unique, common_unique, assume_unique=True)]
    return bool_ind, common_ind[common_inv]

Альтернативная реализация, которая немного (~ 5x) медленнее, но не использует уникальную функцию, находится здесь:

def ismember(a_vec, b_vec):
    ''' MATLAB equivalent ismember function. Slower than above implementation'''
    b_dict = {b_vec[i]: i for i in range(0, len(b_vec))}
    indices = [b_dict.get(x) for x in a_vec if b_dict.get(x) is not None]
    booleans = np.in1d(a_vec, b_vec)
    return booleans, np.array(indices, dtype=int)