从具有 2 维数组的 3 维数组中选择答案

【问题标题】：Select from a 3-dimensional array with a 2-dimensional array从具有 2 维数组的 3 维数组中选择
【发布时间】：2019-05-28 10:27:29
【问题描述】：

我有两个数组：

a: 一个 3 维源数组 (N x M x 2)
b：包含 0 和 1 的二维索引数组 (N x M)。

我想使用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]]
)

起初，我认为这一定有一个简单的解决方案，但我根本找不到。由于我想将它移植到 tensorflow，如果有人知道 numpy 类型的解决方案，我将不胜感激。

编辑：a 的第三个维度可能包含两个以上的元素。因此，b 也可能包含不同于 0 和 1 的索引 - 它不是布尔掩码。

【问题讨论】：

能否详细说明选择的逻辑？我看不出b 和最终结果之间有任何关系。
@ParitoshSingh 对于b 中(m, n) 位置的每个元素x，输出数组在(m, n) 位置应等于a[m, n, x]。

标签： python arrays numpy tensorflow

【解决方案1】：

我们可以为此使用np.where：

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

输出：

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

【讨论】：

我喜欢你在这里使用np.where，但a 的最后一个维度可能包含两个以上的元素——所以我真的需要来自b 的数字索引，而不仅仅是一个布尔掩码。抱歉，我应该在我的问题中澄清这一点。
@mrzo 那么take_along_axis 答案很可能就是你想要的。

【解决方案2】：

正如@jdehesa 所建议的，我们可以使用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]])

【讨论】：

您也可以使用ax0, ax1 = np.ogrid[:b.shape[0], :b.shape[1]] 然后desired = a[ax0, ax1, b]（广播索引）获得相同的结果。
啊，是的，好多了，感谢@jdehesa 的建议

【解决方案3】：

我为 tensorflow 添加了一些解决方案。

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.]]

【讨论】：

不错。作为一个建议，您可以使用idx = tf.stack([row, colum, b], axis=-1) 使idx = tf.concat([row[:,:,tf.newaxis],colum[:,:,tf.newaxis],b[:,:,tf.newaxis]],axis=-1) 更简洁。

【解决方案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]]

【讨论】：