每个元素的索引数组沿二维数组中的第一个维度(numpy.，张量流) [英] Array of indexes for each element alongs the first dimension in a 2D array (numpy., tensorflow)

问题描述

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

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

Now, i would like an array of shape (2,3), where

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

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

What would be the proper way to achieve this? A numpy way that yould generalize to bigger arrays (perhaps even higher dimensional).

请注意与此，其中索引数组包含元组.这不是我要的.

Please note the difference to questions like this, where the array of indexes contains tuples. This is not what I am asking.

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

A more general formulation of the question would be:

• data.shape ==(s0，s1，..，sn)
• indexes.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

推荐答案

以下是NumPy和TensorFlow解决方案:

Here are NumPy and TensorFlow solutions:

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

这篇关于每个元素的索引数组沿二维数组中的第一个维度(numpy.，张量流)的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！