從 numpy 中的索引矩陣創建一個新矩陣

Question

我有一個 3D numpy 數組A形狀（k，l，m）和一個二維 numpy 數組B形狀為（k，l），索引我想要創建一個新的 -1 和 m 之間的特定項目具有形狀 (k,l) 的二維數組C ，如下所示：

import numpy as np
A = np.random.random((2,3,4))
B = np.array([[0,0,0],[2,2,2]))
C = np.zeros((2,3))
for i in range(2):
    for j in range(3):
        C[i,j] = A[i, j, B[i,j]]

有沒有更有效的方法來做到這一點？

Answer 1

使用fromfunction庫函數的內置例程名稱。 並將您的代碼變成

C = np.fromfunction(lambda i, j: A[i, j, B[i,j]], (5, 5))

Answer 2

設置：

import numpy as np
k,l,m = 2,3,4
a = np.arange(k*l*m).reshape(k,l,m)
b = np.random.randint(0,4,(k,l))

print(a)
print('*'*10)
print(b)

---

[[[ 0  1  2  3]
  [ 4  5  6  7]
  [ 8  9 10 11]]

 [[12 13 14 15]
  [16 17 18 19]
  [20 21 22 23]]]
**********
[[3 0 3]
 [2 1 2]]

---

使用 integer 索引到 select 值然后重塑。

x,y = np.indices(a.shape[:-1])
c = a[x,y,b]
print(c)

---

[[ 3  4 11]
 [14 17 22]]

---

使用numpy.ix_ 。

x,y = np.ix_(np.arange(a.shape[0]),np.arange(a.shape[1]))
d = a[x,y,b]