Эффективный поиск экземпляров определенной строки в массиве NumPy

При работе с массивами NumPy можно столкнуться с необходимостью определить, массив содержит определенную строку, но стандартный метод contains для ndarray вызывает вопросы. В этой статье представлены эффективные решения этой проблемы на языке Python.

Один из подходов включает преобразование массива NumPy в список Python с помощью .tolist() и выполнение проверок членства в списке.

a = np.array([[1,2],[10,20],[100,200]])
[1,2] in a.tolist() # Returns True
[1,20] in a.tolist() # Returns False

Другой метод — использовать представление массива и применить функцию .all(1) для сравнения каждой строки с целевой строкой поэлементно .

any((a[:]==[1,2]).all(1)) # Returns True
any((a[:]==[1,20]).all(1)) # Returns False

Кроме того, можно генерировать данные по списку NumPy для потенциального повышения производительности. Однако этот подход может быть неэффективным, если результат не найден раньше.

any(([1,2] == x).all() for x in a) # Stops on first occurrence

Наконец, логические функции NumPy предоставляют краткий способ выполнения сравнений.

any(np.equal(a,[1,2]).all(1)) # Returns True

Результаты тестов показывают, что подпрограммы numpy поддерживают постоянную скорость поиска независимо от сценариев попадания или промаха. Представление, логическое равенство и Python в операторских подходах сопоставимы с точки зрения эффективности, тогда как генератор над NumPy не рекомендуется для поиска по всему массиву.