Python:根据具有三维索引的相应数组将 2D 数组拉伸为 3D

Question

假设我有一些2D数组a = np.ones((3,3))

我想将这个数组拉伸成 3 维。我有数组b，大小与a相同，它提供了 3rd 维索引，a中的每个相应元素也需要去。

我还有充满 NaN 的 3D 阵列c。这是来自a的信息应放入的数组。未被"填充："的剩余空格可以保留 NaN。

>>> a = np.ones((3,3))
>>> b = np.random.randint(0,3,(3,3))
>>> c = np.empty((3,3,3))*np.nan
>>> 
>>> a
array([[ 1.,  1.,  1.],
[ 1.,  1.,  1.],
[ 1.,  1.,  1.]])
>>> b
array([[2, 2, 2],
[1, 0, 2],
[1, 0, 0]])
>>> c
array([[[ nan,  nan,  nan],
[ nan,  nan,  nan],
[ nan,  nan,  nan]],
[[ nan,  nan,  nan],
[ nan,  nan,  nan],
[ nan,  nan,  nan]],
[[ nan,  nan,  nan],
[ nan,  nan,  nan],
[ nan,  nan,  nan]]])

因此，在上面的例子中，我希望以 c[0,0,2] = 1 结束。

我知道我可能会使用一些嵌套循环来做到这一点，但理想情况下，我希望以更有效/矢量化的方式完成此操作。

Answer 1

你可以这样使用花式索引，假设b中的最大值总是小于c.shape[2]：

n1, n2 = a.shape
c[np.arange(n1)[:,None], np.arange(n2), b] = a
c
#array([[[ nan,  nan,   1.],
#        [ nan,  nan,   1.],
#        [ nan,  nan,   1.]],
#       [[ nan,   1.,  nan],
#        [  1.,  nan,  nan],
#        [ nan,  nan,   1.]],
#       [[ nan,   1.,  nan],
#        [  1.,  nan,  nan],
#        [  1.,  nan,  nan]]])

这里

---

我们用所有维度的整数数组来触发高级索引，三个数组相互广播如下(这里我们使用numpy.broacast_arrays来可视化(：

i, j, k = np.broadcast_arrays(np.arange(3)[:,None], np.arange(3), b)
print("first dimension index: ")
print(i)
print("second dimension index: ")
print(j)
print("third dimension index: ")
print(k)
first dimension index: 
[[0 0 0]
[1 1 1]
[2 2 2]]
second dimension index: 
[[0 1 2]
[0 1 2]
[0 1 2]]
third dimension index: 
[[2 2 2]
[1 0 2]
[1 0 0]]

现在，高级索引为(0,0,2(，(0,1,2(，(0,2,2(...，即从每个数组中相同位置选择一个值以形成元素的索引：

---

一些测试案例：

c[0,0,2]
#1.0
c[0,1,2]
#1.0
c[2,1,0]
#1.0

Answer 2

好的，所以这感觉像是一个彻头彻尾的黑客，但可以解决问题：

a = np.ones((3,3))
b = np.array([[2, 2, 2],
[1, 0, 2],
[1, 0, 0]])
c = np.empty((3,3,3))*np.nan
z_coords = np.arange(3)
c[z_coords[None, None, :] == b[..., None]] = a.ravel()

我所做的是创建一个布尔索引数组，该数组适用于我们要分配的索引，然后分配这些索引。

array([[[ nan,  nan,   1.],
[ nan,  nan,   1.],
[ nan,  nan,   1.]],
[[ nan,   1.,  nan],
[  1.,  nan,  nan],
[ nan,  nan,   1.]],
[[ nan,   1.,  nan],
[  1.,  nan,  nan],
[  1.,  nan,  nan]]])

Answer 3

一个更慢但可能更清晰的选项：

x, y = np.indices(c.shape[:2])
c[x, y, b] = a  # same as looping over c[x[i,j], y[i,j], b[i,j]] = a[i,j]

诀窍是生成 3 个具有相同形状的索引数组 - 每个维度一个c。

公认的答案基本上是这样做，但要利用广播