小编典典

每个元素的索引数组沿2D数组中的第一个维度(numpy。,张量流)

python

indexes = np.array([[0,1,3],[1,2,4 ]])
data = np.random.rand(2,5)

现在,我想要一个形状为(2,3)的数组,其中

result[0] = data[0,indexes[0]]
result[1] = data[1,indexes[1]]

实现这一目标的正确方法是什么?一种麻木的方式推广到更大的数组(也许甚至更高的维度)。

请注意区别,以这样的问题这样,在索引的数组包含元组。这不是我要的。

编辑

这个问题的更一般的表述是:

  • data.shape ==(s0,s1,..,sn)
  • indexs.shape ==(s0,s1,…,sn-1,K)
  • 因此,它们具有所有维度,但最后一个相等

result[i, j, ..., k] = data[i, j,...,k, indexes[i, j, ..., k]]

哪里

len([i, j, ..., k]) == len(data)-1 == len(indexes) - 1

阅读 268

收藏
2021-01-20

共1个答案

小编典典

以下是NumPy和TensorFlow解决方案:

import numpy as np
import tensorflow as tf

def gather_index_np(data, index):
    data = np.asarray(data)
    index = np.asarray(index)
    # Make open grid of all but last dimension indices
    grid = np.ogrid[tuple(slice(s) for s in index.shape[:-1])]
    # Add extra dimension in grid
    grid = [g[..., np.newaxis] for g in grid]
    # Complete index
    index_full = tuple(grid + [index])
    # Index data to get result
    result = data[index_full]
    return result

def gather_index_tf(data, index):
    data = tf.convert_to_tensor(data)
    index = tf.convert_to_tensor(index)
    index_shape = tf.shape(index)
    d = index.shape.ndims
    # Make grid of all dimension indices
    grid = tf.meshgrid(*(tf.range(index_shape[i]) for i in range(d)), indexing='ij')
    # Complete index
    index_full = tf.stack(grid[:-1] + [index], axis=-1)
    # Index data to get result
    result = tf.gather_nd(data, index_full)
    return result

例:

import numpy as np
import tensorflow as tf

data = np.arange(10).reshape((2, 5))
index = np.array([[0, 1, 3], [1, 2, 4]])
print(gather_index_np(data, index))
# [[0 1 3]
#  [6 7 9]]
with tf.Session() as sess:
    print(sess.run(gather_index_tf(data, index)))
# [[0 1 3]
#  [6 7 9]]
2021-01-20