带有列表列表的 Numpy 索引[重复]

Question

我有一个numpy 数组的索引列表列表，但在使用它们时并没有完全得到想要的结果。

n = 3
a = np.array([[8, 1, 6],
              [3, 5, 7],
              [4, 9, 2]])
np.random.seed(7)
idx = np.random.choice(np.arange(n), size=(n, n-1))
# array([[0, 1],
#        [2, 0],
#        [1, 2]])

在这种情况下，我想要：

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

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

# 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)

有更简单的方法吗？

Answer 1

你可以在这里使用np.take_along_axis：

np.take_along_axis(a, idx, 1)

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

或者使用broadcasting:

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

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

请注意，您在此处使用integer array indexing，您需要使用idx 指定要在哪些轴和行上建立索引。