Найти элемент по индексу строки и значению

Question

Мне нужно найти индексы массива по индексу строки и значению.

     IN: arr1=[[0   5 ]
               [1   8 ]
               [2   3 ]]

         arr2=[[10  5   10 ]
               [8   10  10 ]
               [3   1   11 ]]
     OUT:
           [[0 1]
            [1 0]
            [2 0]]


     np.argwhere(arr2[arr1[:, 0], :] == arr1[:, 1])

Работает в l oop, но мне нужен более быстрый способ, я пытался сделать это в одном argwhere, но он возвращает пустой массив.

      arr1 = np.array([[0,5],[1,8],[2,3]])

      arr2 = np.array([[10,5,10],[8,10,10],[3,1,11]])

      for (u,v) in arr1:
          result=np.argwhere(arr2[u, :] == v)

Answer 1

Вам просто нужно было нарезать второй столбик arr1, сохраняя размеры. Один из способов сделать это - увеличить dim на np.newaxis/None после нарезки (что уменьшает dims), следовательно -

np.argwhere(arr2[arr1[:, 0], :] == arr1[:, 1, None])

Для компактного вида пропустите индексирование в конце двоеточия -

np.argwhere(arr2[arr1[:,0]] == arr1[:,1,None])

Более компактный способ, но тот, который заставляет копировать второй столбец arr1 -

np.argwhere(arr2[arr1[:,0]] == arr1[:,[1]])

Обратите внимание, что np.argwhere дает нам (row,col) пар в массиве из двух столбцов. С вашим решением на основе l oop вы в основном извлекаете индексы col. Итак, используйте второй столбец вывода из размещенного кода.

Answer 2

Сравнение некоторых характеристик:

%timeit np.argwhere(arr2[arr1[:, 0]] == arr1[:, 1, None])[:, 1]

8,82 мкс ± 531 нс на л oop (среднее ± стандартное отклонение для 7 прогонов, 100000 циклов в каждом)

%timeit [np.argwhere(arr2[u, :] == v)[0,0] for (u,v) in arr1]

20,9 мкс ± 2,74 мкс на л oop (среднее ± стандартное отклонение для 7 прогонов, 10000 циклов в каждом)