有效地定义一个隐式 Numpy 数组

Question

A 和 B 是常见形状 [n1,n2,n3] 的 Numpy 数组。 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" 的答案中。这里唯一真正新的是沿 3 维数组的一个轴进行反转。