Найти совпадающие строки в 2-мерной матрице

Я хотел бы получить индекс 2-мерного массива Numpy, который соответствует строке. Например, мой массив:

vals = np.array([[0, 0],
                 [1, 0],
                 [2, 0],
                 [0, 1],
                 [1, 1],
                 [2, 1],
                 [0, 2],
                 [1, 2],
                 [2, 2],
                 [0, 3],
                 [1, 3],
                 [2, 3],
                 [0, 0],
                 [1, 0],
                 [2, 0],
                 [0, 1],
                 [1, 1],
                 [2, 1],
                 [0, 2],
                 [1, 2],
                 [2, 2],
                 [0, 3],
                 [1, 3],
                 [2, 3]])

Я хотел бы получить индекс, который соответствует строке [0, 1], которая является индексом 3 и 15. Когда я делаю что-то вроде numpy.where(vals == [0 ,1]), я получаю...

(array([ 0,  3,  3,  4,  5,  6,  9, 12, 15, 15, 16, 17, 18, 21]), array([0, 0, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 0, 0]))

Я хочу индексный массив ([3, 15]).

Ответ 1

Вам нужна функция np.where, чтобы получить индексы:

>>> np.where((vals == (0, 1)).all(axis=1))
(array([ 3, 15]),)

Чтобы разобрать это:

>>> vals == (0, 1)
array([[ True, False],
       [False, False],
       ...
       [ True, False],
       [False, False],
       [False, False]], dtype=bool)

и вызов метода .all в этом массиве (с axis=1) дает вам True, где оба значения True:

>>> (vals == (0, 1)).all(axis=1)
array([False, False, False,  True, False, False, False, False, False,
       False, False, False, False, False, False,  True, False, False,
       False, False, False, False, False, False], dtype=bool)

и получить, какие индексы True:

>>> np.where((vals == (0, 1)).all(axis=1))
(array([ 3, 15]),)

Я считаю, что мое решение немного читаемо, но, как указывает unutbu, следующее может быть быстрее и возвращает то же значение, что и (vals == (0, 1)).all(axis=1):

>>> (vals[:, 0] == 0) & (vals[:, 1] == 1)

Ответ 2

In [5]: np.where((vals[:,0] == 0) & (vals[:,1]==1))[0]
Out[5]: array([ 3, 15])

Я не уверен, почему, но это значительно быстрее, чем np.where((vals == (0, 1)).all(axis=1)):

In [34]: vals2 = np.tile(vals, (1000,1))

In [35]: %timeit np.where((vals2 == (0, 1)).all(axis=1))[0]
1000 loops, best of 3: 808 µs per loop

In [36]: %timeit np.where((vals2[:,0] == 0) & (vals2[:,1]==1))[0]
10000 loops, best of 3: 152 µs per loop

Ответ 3

Используя пакет numpy_indexed, вы можете просто написать:

import numpy_indexed as npi
print(np.flatnonzero(npi.contains([[0, 1]], vals)))

Ответ 4

Мне было интересно, как эффективно найти индексы строк в двумерном массиве Numpy, которые соответствуют строкам из другого 2-мерного массива Numpy. Точнее, учитывая vals, мне нужны индексы, которые соответствуют строкам a = np.array([[1, 0], [0,0], [0,1]]), которые являются индексами 0,1,3,12,13,15. Конечно, это можно сделать, как и раньше, используя цикл над строками в a, но мне было интересно, есть ли способ без использования цикла?