Question

A和B是常见形状{N}的Numpy数组[n1,n2,n3]。 B的值都是[0,n3)中的整数。我希望A“反转”B，因为A的每个值都满足A[i,j,B[i,j,k]]=k在适当范围内的所有i,j,k。用for循环做这个是多么明显，但我怀疑有一个聪明的单行使用花式索引。有人看到了吗？

Answer 1

以下是两种方法。

第一种方法是单行：A = B.argsort(axis=-1)

这是一个例子。 B具有形状（3,5,7），对于每个固定的i和j，B[i,j,:]是range(B.shape[2])的排列。

In [386]: B
Out[386]: 
array([[[1, 5, 4, 6, 2, 3, 0],
        [6, 5, 3, 4, 2, 1, 0],
        [4, 5, 0, 3, 1, 2, 6],
        [0, 5, 6, 3, 2, 1, 4],
        [4, 1, 5, 2, 6, 3, 0]],

       [[2, 6, 0, 1, 5, 4, 3],
        [3, 2, 4, 0, 1, 5, 6],
        [3, 4, 6, 5, 1, 2, 0],
        [4, 6, 3, 0, 2, 5, 1],
        [0, 3, 1, 6, 4, 5, 2]],

       [[0, 3, 6, 2, 1, 5, 4],
        [3, 1, 2, 4, 6, 0, 5],
        [1, 3, 5, 6, 4, 0, 2],
        [4, 1, 6, 0, 2, 3, 5],
        [6, 4, 5, 1, 0, 3, 2]]])

In [387]: A = B.argsort(axis=-1)

In [388]: A
Out[388]: 
array([[[6, 0, 4, 5, 2, 1, 3],
        [6, 5, 4, 2, 3, 1, 0],
        [2, 4, 5, 3, 0, 1, 6],
        [0, 5, 4, 3, 6, 1, 2],
        [6, 1, 3, 5, 0, 2, 4]],

       [[2, 3, 0, 6, 5, 4, 1],
        [3, 4, 1, 0, 2, 5, 6],
        [6, 4, 5, 0, 1, 3, 2],
        [3, 6, 4, 2, 0, 5, 1],
        [0, 2, 6, 1, 4, 5, 3]],

       [[0, 4, 3, 1, 6, 5, 2],
        [5, 1, 2, 0, 3, 6, 4],
        [5, 0, 6, 1, 4, 2, 3],
        [3, 1, 4, 5, 0, 6, 2],
        [4, 3, 6, 5, 1, 2, 0]]])

通过抽样几个值来验证所需的属性。

In [389]: A[0, 0, B[0, 0, 0]]
Out[389]: 0

In [390]: A[0, 0, B[0, 0, 1]]
Out[390]: 1

In [391]: A[0, 0, B[0, 0, :]]
Out[391]: array([0, 1, 2, 3, 4, 5, 6])

In [392]: A[2, 3, B[2, 3, :]]
Out[392]: array([0, 1, 2, 3, 4, 5, 6])

第二种方法比使用argsort具有更低的时间复杂度，但它是三线而不是一线。我会使用与上面相同的B。

创建A，但尚未分配任何值。

In [393]: A = np.empty_like(B)

为B的每个维度创建索引数组。

In [394]: i, j, k = np.ogrid[[slice(n) for n in B.shape]]  # or np.ix_(*[range(n) for n in B.shape])

这是很酷的部分。完全按照你在问题中写的那样完成作业。

In [395]: A[i, j, B[i, j, k]] = k

确认我们与上述A相同。

In [396]: A
Out[396]: 
array([[[6, 0, 4, 5, 2, 1, 3],
        [6, 5, 4, 2, 3, 1, 0],
        [2, 4, 5, 3, 0, 1, 6],
        [0, 5, 4, 3, 6, 1, 2],
        [6, 1, 3, 5, 0, 2, 4]],

       [[2, 3, 0, 6, 5, 4, 1],
        [3, 4, 1, 0, 2, 5, 6],
        [6, 4, 5, 0, 1, 3, 2],
        [3, 6, 4, 2, 0, 5, 1],
        [0, 2, 6, 1, 4, 5, 3]],

       [[0, 4, 3, 1, 6, 5, 2],
        [5, 1, 2, 0, 3, 6, 4],
        [5, 0, 6, 1, 4, 2, 3],
        [3, 1, 4, 5, 0, 6, 2],
        [4, 3, 6, 5, 1, 2, 0]]])

在SO上搜索更多内容后，我发现这两种方法都出现在问题"How to invert a permutation array in numpy"的答案中。这里唯一真正新的是沿着三维阵列的一个轴进行反演。

有效地定义隐式Numpy数组

1 个答案: