Python Numpy находит индекс строки массива, состоящего из 5 значений, путем поиска 2 значений

Question

Привет, извиняюсь, если название сбивает с толку, я новичок в numpy и не привык к терминологии.

Предположим, у нас есть массив numpy, действующий как карта мира. Параметры (x y r g b) - все являются int16

Пример:

a = np.array([[  0,   0,   0, 255,   0], #index 0
              [  0,   1,   0,   0, 255], #index 1
              [  0,   2,   0, 255,   0]]) #index 2

Теперь мы хотим найти значение индекса строки со значениями x и y (0, 2) - отсюда и строка с индексом 2.

[  0,   1,   0,   0, 255] #index 2

Как бы я это сделал, не вводя остальные значения (r g b)? В основном мы ищем строку из пяти значений с двумя значениями - как бы я это сделал?

Answer 1

Вы можете разрезать строки до второго столбца и проверить, равны ли они [0,2].Затем используйте all, установив axis в 1, чтобы установить в True те, которые удовлетворяют всем условиям, и используйте логический массив для индексации ndarray:

a = np.array([[  0,   0,   0, 255,   0],
              [  0,   1,   0,   0, 255],
              [  0,   2,   0, 255,   0]])

a[(a[:,:2] == [0,2]).all(1)]
# array([[  0,   2,   0, 255,   0]])

Answer 2

Вот ваши данные:

import numpy as np

arr = np.array([[  0,   0,   0, 255,   0],
              [  0,   1,   0,   0, 255],
              [  0,   2,   0, 255,   0]])

a,b = 0,2 # [a,b] is what we are looking for, in the first two cols

Вот решение для получения индекса строки, содержащего [a, b]:

found_index = np.argmax(np.logical_and(arr[:,0]==[a],arr[:,1]==[b]))
print (found_index)

Выход:

2

Пояснение:

Лучший способ понять, как это работает, - напечатать каждую его часть:

print (arr[:,0]==[a])

Выходы:

[True True True]

print (arr[:,1]==[b])

Выходы:

[False False True]

print (np.logical_and(arr[:,0]==[a],arr[:,1]==[b]))
# print (np.logical_and([ True  True  True], [False False  True]))

Выходы:

[Ложь Ложь Истина]