Массив индексов для каждого элемента вдоль первого измерения в двумерном массиве (numpy., Tenorflow)

Question

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

Каким был бы правильный способ добиться этого?Обалденный способ, который вы можете обобщить на большие массивы (возможно, даже на более высокий размер).

Обратите внимание на разницу с такими вопросами, как this , где массив индексов содержит кортежи.Это не то, что я спрашиваю.

Редактировать

Более общая формулировка вопроса была бы такой:

• 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

Answer 1

numpy имеет take_along_axis, что делает то, что вы описываете, плюс оно также позволяет вам выбрать ось.

Пример:

>>> a = np.arange(24).reshape(2,3,4)
>>> i = np.random.randint(0,4,(2,3,5))
>>> i
array([[[3, 3, 0, 1, 3],
        [3, 1, 0, 3, 3],
        [3, 2, 0, 3, 3]],

       [[2, 3, 0, 0, 0],
        [1, 1, 3, 1, 2],
        [1, 3, 0, 0, 2]]])

>>> np.take_along_axis(a, i, -1)
array([[[ 3,  3,  0,  1,  3],
        [ 7,  5,  4,  7,  7],
        [11, 10,  8, 11, 11]],

       [[14, 15, 12, 12, 12],
        [17, 17, 19, 17, 18],
        [21, 23, 20, 20, 22]]])

Answer 2

Вот решения 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]]