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]]
实现这一目标的正确方法是什么?一种麻木的方式推广到更大的数组(也许甚至更高的维度)。
请注意区别,以这样的问题这样,在索引的数组包含元组。这不是我要的。
这个问题的更一般的表述是:
比
result[i, j, ..., k] = data[i, j,...,k, indexes[i, j, ..., k]]
哪里
len([i, j, ..., k]) == len(data)-1 == len(indexes) - 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]]