Numpy индексный вложенный массив с индексами из другого массива

Question

У меня есть массив numpy, например:

u = np.arange(10).reshape(5, 2)

array([[0, 1],
   [2, 3],
   [4, 5],
   [6, 7],
   [8, 9]])

У меня есть второй массив, например

a = np.array([1,0,0,1,0])

Я хотел бы использовать значения из a для индексации подмассивов и. E.g. a[0] is 1, so we chose u[0,1], a[1] is 0, so we choose u[1, 0] и т. Д.

Я много чего перепробовал и хотел бы сделать это без циклов. Даже после прочтения руководства по индексированию numpys я так и не нашел, как это сделать.

То, что я пробовал, не удалось:

>>> u[:, [0,0,1,0,1]]
array([[0, 0, 1, 0, 1],
   [2, 2, 3, 2, 3],
   [4, 4, 5, 4, 5],
   [6, 6, 7, 6, 7],
   [8, 8, 9, 8, 9]])

u[[True, False, True, True, True]]
array([[0, 1],
   [4, 5],
   [6, 7],
   [8, 9]])

Наконец, чтобы устранить путаницу вот что я хочу, однако с python петлями:

>>> x = []
>>> ct = 0
>>> for i in u:
        x.append(i[a[ct]])
        ct += 1

>>> x
[1, 2, 4, 7, 8]

Заранее спасибо.

Answer 1

Использование:

import numpy as np

u = np.arange(10).reshape(5, 2)
a = np.array([1, 0, 0, 1, 0])
r, _ = u.shape  # get how many rows to use in np.arange(r)

print(u[np.arange(r), a])

Вывод

[1 2 4 7 8]

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

Answer 2

Рекомендуется использовать словари . Вы сможете получить доступ к элементу своего словаря, используя указанный ключ c (в вашем примере это просто число). Это очень мощный инструмент. Конечно, словари - это гораздо больше.

Я действительно верю в c мощь StackOverflow, блогов и целых rnet педагогов в целом. Возникшая проблема - это правильная возможность узнать что-то новое, мощное и полезное во многих областях. Полный обзор словаря займет всего 30 минут. Стоит ли тестироваться?