Python 使用列表列表进行Numpy索引

我有一个

numpy

数组的索引列表，但使用它们时并没有得到想要的结果

n=3
a=np.array（[[8,1,6]，
[3, 5, 7],
[4, 9, 2]])
np.随机种子（7）
idx=np.random.choice（np.arange（n），size=（n，n-1））
#数组（[[0，1]，
#        [2, 0],
#        [1, 2]])

在这种情况下，我希望：

• 第0行的元素0和1
• 第1行的元素2和0
• 第2行的元素1和2

我的列表有

n个子列表

，所有这些列表的长度都相同。 我希望每个子列表只使用一次，而不是用于所有轴

# Wanted result
# b = array[[8, 1],
#           [7, 3],
#           [9, 2]])

我可以做到这一点，但这似乎是相当繁琐的大量重复和重塑

# Possibility 1
b = a[:, idx]
# array([[[8, 1],   | [[3, 5],   |  [[4, 9],
#         [6, 8],   |  [7, 3],   |   [2, 4],
#         [1, 6]],  |  [5, 7]],  |   [9, 2]])
b = b[np.arange(n), np.arange(n), :]

# Possibility 2
b = a[np.repeat(range(n), n-1), idx.ravel()]
# array([8, 1, 7, 3, 9, 2])
b = b.reshape(n, n-1)

---

有更简单的方法吗？

您可以使用

np。沿轴取下\u

np.take_along_axis(a, idx, 1)

array([[8, 1],
       [7, 3],
       [9, 2]])

或使用
广播
：

a[np.arange(a.shape[0])[:,None], idx]

array([[8, 1],
       [7, 3],
       [9, 2]])

注意，在这里使用时，需要指定要使用
idx
索引的轴和行