Numpy - использовать значения в качестве индекса других массивов

Question

У меня проблема с Numpy, я пытаюсь использовать значения каждой строки (из B) в качестве индексов другого многомерного массива (A):

>>> A
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])
>>> B
    array([[ 0,  1,  2,  3],
           [ 1,  5,  6,  7],
           [ 8,  9, 10, 11]])
>>> np.clip(B, 0, 2)
    array([[0, 1, 2, 2],
           [1, 2, 2, 2],
           [2, 2, 2, 2]])

Вот ожидаемый результат :

array([[0, 1, 2, 2],
       [4, 5, 5, 5],
       [8, 8, 8, 8]])

Есть идеи? Большое спасибо за вашу помощь.

Answer 1

Похоже, что вы ищете np.take_along_axis:

A = np.array([
       [0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]]
)
B = np.array([
           [ 0,  1,  2,  3],
           [ 1,  5,  6,  7],
           [ 8,  9, 10, 11]]
)
C = np.clip(B, 0, 2)

res = np.take_along_axis(A, C, 1)
print(res)

Answer 2

Попробуйте:

result = np.clip(B, 0, 2)
for i, row in enumerate(result):
    for j, elem in enumerate(row):
        result[i][j] = A[i][elem]

После ответа @ hilberts_drinking_problem я провел несколько тестов с использованием timeit. Ответы при тестировании с 10,000,000 loops заняли 103s и 82s, последний был моим. Поэтому, если ваши данные невелики, но если вы собираетесь зацикливаться на них, я бы порекомендовал вам использовать мой ответ. Если это не так, я считаю, что удобство ответа hilberts_drinking_problem делает его превосходным. Для получения дополнительной информации см. Комментарии к этому ответу.

Обратите внимание, что в обоих случаях я протестировал только детали после отсечения, то есть предыдущая подготовительная деталь не влияет на эталонный тест.