Выберите из 3-мерного массива с 2-мерным массивом

Question

У меня есть два массива:

• a: трехмерный исходный массив ( N x M x 2 )
• b:массив двумерных индексов ( N x M ), содержащий 0 и 1 с.

Я хочу использовать индексы в b для выбора соответствующих элементов aв своем третьем измерении.Полученный массив должен иметь размеры N x M .Вот пример кода:

import numpy as np

a = np.array( # dims: 3x3x2
    [[[ 0,  1],
     [ 2,  3],
     [ 4,  5]],
    [[ 6,  7],
     [ 8,  9],
     [10, 11]],
    [[12, 13],
     [14, 15],
     [16, 17]]]
)
b = np.array( # dims: 3x3
    [[1, 1, 1],
    [1, 1, 1],
    [1, 1, 1]]
)

# select the elements in a according to b
# to achieve this result:
desired = np.array(
  [[ 1,  3,  5],
   [ 7,  9, 11],
   [13, 15, 17]]
)

Сначала я подумал, что у этого должно быть простое решение, но я не смог его найти вообще.Так как я хотел бы перенести его на tenorflow, я был бы признателен, если бы кто-нибудь знал для этого решение типа numpy.

Edit: Третье измерение a может содержать более двухэлементы.Следовательно, b может также содержать индексы, отличные от 0 и 1 - это не логическая маска.

Answer 1

Я добавил несколько решений для тензорного потока.

import tensorflow as tf

a = tf.constant([[[ 0,  1],[ 2,  3],[ 4,  5]],
                 [[ 6,  7],[ 8,  9],[10, 11]],
                 [[12, 13],[14, 15],[16, 17]]],dtype=tf.float32)
b = tf.constant([[1, 1, 1],[1, 1, 1],[1, 1, 1]],dtype=tf.int32)

# 1. use tf.gather_nd
colum,row = tf.meshgrid(tf.range(a.shape[0]),tf.range(a.shape[1]))
idx = tf.stack([row, colum, b], axis=-1) # Thanks for @jdehesa's suggestion
result1 = tf.gather_nd(a,idx)

# 2. use tf.reduce_sum
mask = tf.one_hot(b,depth=a.shape[-1],dtype=tf.float32)
result2 = tf.reduce_sum(a*mask,axis=-1)

# 3. use tf.boolean_mask
mask = tf.one_hot(b,depth=a.shape[-1],dtype=tf.float32)
result3 = tf.reshape(tf.boolean_mask(a,mask),b.shape)

with tf.Session() as sess:
    print('method 1: \n',sess.run(result1))
    print('method 2: \n',sess.run(result2))
    print('method 3: \n',sess.run(result3))

method 1: 
 [[ 1.  3.  5.]
 [ 7.  9. 11.]
 [13. 15. 17.]]
method 2: 
 [[ 1.  3.  5.]
 [ 7.  9. 11.]
 [13. 15. 17.]]
method 3: 
 [[ 1.  3.  5.]
 [ 7.  9. 11.]
 [13. 15. 17.]]

Answer 2

Мы можем использовать np.where для этого:

np.where(b, a[:, :, 1], a[:, :, 0])

Выход:

array([[ 1,  3,  5],
       [ 7,  9, 11],
       [13, 15, 17]])

Answer 3

Как @jdehesa sugests, мы можем использовать np.ogrid, чтобы получить индексы для первых двух осей:

ax0, ax1 = np.ogrid[:b.shape[0], :b.shape[1]]

И затем мы можем использовать b для прямого индексированиявдоль последней оси.Обратите внимание, что ax0 и ax1 будут транслироваться в формате b:

desired = a[ax0, ax1 ,b] 

print(desired)
array([[ 1,  3,  5],
       [ 7,  9, 11],
       [13, 15, 17]])

Answer 4

Вы можете использовать np.take_along_axis:

import numpy as np

a = np.array(
    [[[ 0,  1],
      [ 2,  3],
      [ 4,  5]],
     [[ 6,  7],
      [ 8,  9],
      [10, 11]],
     [[12, 13],
      [14, 15],
      [16, 17]]])
b = np.array(
    [[1, 1, 1],
     [1, 1, 1],
     [1, 1, 1]])
print(np.take_along_axis(a, b[..., np.newaxis], axis=-1)[..., 0])
# [[ 1  3  5]
#  [ 7  9 11]
#  [13 15 17]]