Question

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

我想将此数组拉伸为3维。我的数组b与a的大小相同，它提供了第3维中的索引，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]]])

这里我们使用所有维度的整数数组来触发advanced indexing，并且三个数组相互广播如下（这里我们使用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一个。

接受的答案基本上是这样做，但利用广播

Python：基于具有第三维索引

3 个答案: