由多维数组的索引矩阵指定的Numpy选择矩阵

时间:2019-01-27 23:52:35

标签: python arrays numpy indexing

我有一个大小为a的numpy数组5x5x4x5x5。我还有另一个大小为b的矩阵5x5。我想将a[i,j,b[i,j]]的{​​{1}}从0变为4,将i的{​​{1}}从0变为4。这将给我一个j矩阵。有没有办法只使用两个5x5x1x5x5循环吗?

2 个答案:

答案 0 :(得分:1)

让我们将矩阵a视为大小为(= 5 x 5 x 4)的100个(5, 5)矩阵。因此,如果您可以获得每个三元组的衬里索引-(i, j, b[i, j])-您就完成了。这就是np.ravel_multi_index出现的地方。以下是代码。

import numpy as np
import itertools

# create some matrices
a = np.random.randint(0, 10, (5, 5, 4, 5, 5))
b = np.random(0, 4, (5, 5))

# creating all possible triplets - (ind1, ind2, ind3)
inds = list(itertools.product(range(5), range(5)))
(ind1, ind2), ind3 = zip(*inds), b.flatten()

allInds = np.array([ind1, ind2, ind3])
linearInds = np.ravel_multi_index(allInds, (5,5,4))

# reshaping the input array
a_reshaped = np.reshape(a, (100, 5, 5))

# selecting the appropriate indices
res1 = a_reshaped[linearInds, :, :]

# reshaping back into desired shape
res1 = np.reshape(res1, (5, 5, 1, 5, 5))

# verifying with the brute force method
res2 = np.empty((5, 5, 1, 5, 5))
for i in range(5):
    for j in range(5):
        res2[i, j, 0] = a[i, j, b[i, j], :, :]

print np.all(res1 == res2)  # should print True

答案 1 :(得分:0)

np.take_along_axis正是出于这个目的-

   Text                 Label
1. I like to travel     Travel
2. I love to travel     Unmapped