如何使用 2d 数组的值作为 numpy 中 3d 数组的索引？

Question

我有一个二维数组 A，其形状为 (n x m)，其中位置 (i,j) 处的数组的每个元素都包含第三个值 k。我想根据 2d 数组值和位置在位置 (k,i,j) 增加一个维度为 nxmxl 的 3d 数组。

例如如果

A = [[0,1],[3,3]] -> I would want B to be 

[[[1,0],
  [0,0]],

  [0,1],
  [0,0]],

  [0,0],
  [0,1]],

  [0,0],
  [0,2]]]

如何在 numpy 中高效地执行此操作？

Answer 1

问题有点模棱两可，但如果目的是在索引 (0,0,0)、(1,0,1)、(3,1,0) 和 [=17 处递增某个未知数组 B =]，那么下面的应该没问题：

B[(A.ravel(), ) + np.unravel_index(np.arange(np.prod(A.shape)), A.shape)] += increment

例如：

A = np.array([[0,1],[3,3]])
B = np.zeros((4,2,2), dtype=int)
increment = 1

B[(A.ravel(), ) + np.unravel_index(np.arange(np.prod(A.shape)), A.shape)] += increment

>>> B
array([[[1, 0],
        [0, 0]],

       [[0, 1],
        [0, 0]],

       [[0, 0],
        [0, 0]],

       [[0, 0],
        [1, 1]]])

做同样事情的另一种方法是：

w, h = A.shape
indices = (A.ravel(),) + tuple(np.mgrid[:w, :h].reshape(2, -1))

# then
B[indices] += increment

Answer 2

我可以生产你的B：

In [208]: res = np.zeros((4,2,2),int)
In [209]: res.reshape(4,4)[np.arange(4), A.ravel()] = [1,1,1,2]
In [210]: res
Out[210]: 
array([[[1, 0],
        [0, 0]],

       [[0, 1],
        [0, 0]],

       [[0, 0],
        [0, 1]],

       [[0, 0],
        [0, 2]]])

我使用 reshape 因为 A 值看起来更像

的索引

In [211]: res.reshape(4,4)
Out[211]: 
array([[1, 0, 0, 0],
       [0, 1, 0, 0],
       [0, 0, 0, 1],
       [0, 0, 0, 2]])

如何使用 2d 数组的值作为 numpy 中 3d 数组的索引？

How to use the value of a 2d array as an index to a 3d array in numpy?

python

arrays

indexing

numpy