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 的答案的便利性使其更胜一筹。有关详细信息，请参阅此答案的 cmets。

请注意，在这两种情况下，我只测试了裁剪后的部分，这意味着前面的准备部分不会干扰基准测试。