найти массив NumPy в другом массиве NUMPY

Question

Мне нужно найти маленький массив NumPy в гораздо большем массиве NUMPY.Например:

import numpy as np
a = np.array([1, 1])
b = np.array([2, 3, 3, 1, 1, 1, 8, 3, 1, 6, 0, 1, 1, 3, 4])

Функция

find_numpy_array_in_other_numpy_array(a, b)

должна возвращать индексы

[3, 4, 11]

, которые представляют место, где полный массив NumPy a появляется в полной NUMPYмассив b.

Существует грубый подход к этой проблеме, который медленен при работе с очень большими b массивами:

ok = []
for idx in range(b.size - a.size + 1):
    if np.all(a == b[idx : idx + a.size]):
        ok.append(idx)

Я ищу гораздо более быстрый путьнайти все индексы полного массива a в массиве b.Быстрый подход также должен позволять другие функции сравнения, например, для нахождения наихудшей разницы между a и b:

diffs = []
for idx in range(b.size - a.size + 1):
    bi = b[idx : idx + a.size]
    diff = np.nanmax(np.abs(bi - a))
    diffs.append(diff)

Answer 1

Следующий код находит все совпадения 1-го элемента в вашей последовательности (a) в массиве b.Затем он создает новый массив со столбцами возможных кандидатов на последовательность, сравнивает их с полной последовательностью и фильтрует исходные индексы

seq, arr = a, b
len_seq = len(seq)    
ini_idx = (arr[:-len_seq+1]==seq[0]).nonzero()[0] # idx of possible sequence canditates   
seq_candidates = arr[np.arange(1, len_seq)[:, None]+ini_idx] # columns with possible seq. candidates   
mask = (seq_candidates==seq[1:,None]).all(axis=0)
idx = ini_idx[mask]

Answer 2

Вы можете использовать Numba для компиляции функции.Вы можете сделать это так:

import numpy as np
import numba as nb

@nb.njit(parallel=True)
def search_in_array(a, b):
    idx = np.empty(len(b) - len(a) + 1, dtype=np.bool_)
    for i in nb.prange(len(idx)):
        idx[i] = np.all(a == b[i:i + len(a)])
    return np.where(idx)[0]

a = np.array([1, 1])
b = np.array([2, 3, 3, 1, 1, 1, 8, 3, 1, 6, 0, 1, 1, 3, 4])
print(search_in_array(a, b))
# [ 3  4 11]

Быстрый тест:

import numpy as np

np.random.seed(100)
a = np.random.randint(5, size=10)
b = np.random.randint(5, size=10_000_000)

# Non-compiled function
%timeit search_in_array.py_func(a, b)
# 22.8 s ± 242 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# Compiled function
%timeit search_in_array(a, b)
# 54.7 ms ± 1.31 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

Как видите, вы можете получить ~ 400-кратное ускорение и стоимость памяти относительно низка (логическое значениемассив того же размера, что и большой массив).

Answer 3

Настройка общего решения

Для общего решения мы можем создать 2D массив скользящих окон и затем выполнить соответствующие операции -

from skimage.util.shape import view_as_windows

b2D = view_as_windows(b,len(a))

NumPy equivalent implementation.

Задача № 1

Затем для решения проблемы соответствия индексов достаточно просто -

matching_indices = np.flatnonzero((b2D==a).all(axis=1))

Задача #2

Чтобы решить вторую проблему, она легко сопоставляется с учетом того, что любая операция сокращения уфунса для получения выходного элемента должна переводиться в уменьшение по второй оси в предлагаемом решении с использованием этогоaxis аргумент ufunc -

diffs = np.nanmax(np.abs(b2D-a),axis=1)