Блог

Главная - Вопросы по программированию - Как найти индексы строк в A, где столбцы 1 и 3 равны B (без цикла for) в numpy?

Как найти индексы строк в A, где столбцы 1 и 3 равны B (без цикла for) в numpy?

#python #numpy

Вопрос:

Даны 2 массива numpy A и B, как показано ниже :

 A =
[[  1, 200,  15],
 [  0, 600,  25],
 [  2, 200,  20],
 [  3, 100,  10],
 [  1, 300,   5],
 [  4, 400,   3],
 [  0, 100,  12],
 [  5, 300,   2],
 [  2, 300,  25],
 [  6, 100,   1],
 [  1, 400,  74],
 [  7, 300,  10],
 [  2, 400,  15],
 [  8, 200,  13],
 [  3, 400,  29],
 [  3, 600,  37],
 [  0, 200,  20],
 [  9, 300,  42],
 [  5, 400,  30],
 [  6, 200,  51],
 [  7, 400,  21],
 [  9, 500,   5],
 [ 10, 300,   9] ]


B =
[[ 0, 25.],
 [ 1, 74.],
 [ 2, 25.],
 [ 3, 37.],
 [ 4,  3.],
 [ 5, 30.],
 [ 6, 51.],
 [ 7, 21.],
 [ 8, 13.],
 [ 9, 42.],
 [10,  9.]]

Каков эффективный способ (без цикла FOR) вычисления индексов строк в массиве A
, где столбцы 0 и 2 A образуют массив B?

Ожидаемый ответ-я:

 i = [ 1,  5,  8, 10, 13, 15, 17, 18, 19, 20, 22]

Так что:

 A[i,:] == B

Заранее благодарю вас

Ответ №1:

Вы могли бы сделать следующее:

  1. Начните с захвата первого и третьего столбцов A с 
     >>> A[:,(0,2)]
  2. Затем сравните все возможные пары строк (A[i], B[j]) с помощью broadcasting трюка и np.equal : 
     >>> np.equal(A[:,(0,2)][:,None], B[None])
  3. Убедитесь, что все элементы в строке совпадают, используя np.all : 
     >>> np.equal(A[:,(0,2)][:,None], B[None]).all(2)
  4. Наконец, посмотрите на argmax на первой оси, это предполагает, что есть по крайней мере одно совпадение A для каждой строки B : 
     >>> np.equal(A[:,(0,2)][:,None], B[None]).all(2).argmax(0)

    Приведенная выше строка не будет работать, например, если вы удалите B последнюю строку. В одной из строк все равные условия являются ложными (вообще не совпадают B ), но argmax все равно выберет первое False значение, т. е. возвращающий индекс 0 … что не является очень общим решением.

    Более точный метод состоит в том, чтобы np.nonzero вместо этого использовать только ненулевые значения, что именно то, что мы здесь хотим:

     >>> np.equal(A[:,(0,2)][:,None], B[None]).all(2).nonzero()[0]

Комментарии:

1. Можете ли вы немного объяснить «трюк с трансляцией»? Я видел это несколько раз, но мне трудно представить, что именно здесь происходит. Спасибо.

2. Этот трюк включает в себя добавление дополнительного измерения с противоположных сторон. Возьмем этот пример: A = np.arange(1,4) и B = np.arange(1,4) . Тогда A[None] будет иметь форму (1, 3) , в то время B[:, None] как будет иметь форму (3, 1) . При использовании в сочетании с оператором (выше это было np.equal ), тогда операция не будет выполняться по элементам (это было бы в случае , если бы вы дали A , и B ), а вместо этого будут рассмотрены все комбинации пар. Наглядным примером является * оператор: A*B является ли элементарный продукт, в то время A[None]*B[:,None] как внешний продукт между A и B .

3. Огромное спасибо. Думаю, теперь он у меня есть. Я ценю то время, что вы добровольно согласились на ЭТО. Будьте в безопасности и оставайтесь здоровыми.

4. Спасибо, Иван, это работает…

Метки: ,