python - Numpy 中的类似分散操作

Question

我面临以下问题:

我有一个形状为 (*S, N, N) 的 numpy 数组 A，其中 S 是任意元组， N 是某个正整数。另一个形状为 (*S, 2) 的数组 I 表示 A 中的索引。 I 中的值是 {0, ..., N-1} 中的整数。

我要设置

A[i1, ...., ik, I[i1, ..., ik, 0], I[i1, ...., ik, 1]] := 0(伪代码)

对于所有有效索引i1, ..., ik。

例如，考虑以下示例，其中 N=3 和 S=(2):

A = np.array([[[ 1,  2,  3],
               [ 4,  5,  6],
               [ 7,  8,  9]],

              [[10, 11, 12],
               [13, 14, 15],
               [16, 17, 18]]])

I = np.array([[0, 2],
              [2, 1]])

在这种情况下所需的输出是

np.array([[[ 1,  2,  0],
           [ 4,  5,  6],
           [ 7,  8,  9]],

          [[10, 11, 12],
           [13, 14, 15],
           [16, 0, 18]]])

请注意，索引 (0, 0, 2) 处的 3(对应于值 I[0] 为 np.array([0, 2])) 和索引 (1, 2, 1) 处的 17(对应于 I[1] 的值 np.array ([2, 1])) 更改为 0。

这看起来有点像 PyTorch 中的分散操作。然而，它并不完全符合这个问题。此外，numpy 不提供分散方法。我已经阅读了 numpy 关于索引和各种分配值的方法的文档。尽管如此，我还没有想出一种聪明的方法来解决这个问题(即没有循环的方法)。

如果您有建议，我将不胜感激!

Answer 1

有趣的问题!

我相信您可以简单地将 A 和 I reshape 为 (-1,N,N) 和 (-1, 2)分别并始终处理(K x N x N)中的问题，然后简单地 reshape 回到原始尺寸。我建议的解决方案如下:

import numpy as np

# Just extend the problem a little because why not
A = np.array([[[ 1,  2,  3],
               [ 4,  5,  6],
               [ 7,  8,  9]],

              [[10, 11, 12],
               [13, 14, 15],
               [16, 17, 18]],
            
              [[19, 20, 21],
               [22, 23, 24],
               [25, 26, 27]],
])

I = np.array([[0, 2],
              [2, 1],
              [1, 1]])


B = np.stack([A,A])
C = np.array([
    [[0,2],
     [2,1],
     [1,1]],
    
    [[0,0],
     [1,1],
     [2,2]],
])

def scatterlike(A: np.ndarray, I: np.ndarray, target: float=0):
    A_ = A.copy().reshape(-1, *A.shape[-2:])
    A_[(range(len(A_)),*I.reshape(-1,2).T.tolist())] = target
    return A_.reshape(A.shape)

我在代码中做了一个A.copy()，因为它更容易调试，但如果您愿意，您可以就地完成所有操作。我很确定这段代码也适用于 (*S, N, M) ，这很好。我避免使用 np.ndarray 类型进行索引，因为 np.ndarray 和其他可迭代对象的索引行为不同。

以下是一些输出:

# Simplest case  
scatterlike(A[0], I[0])
array([[1, 2, 0],
       [4, 5, 6],
       [7, 8, 9]])

# Extended version of the example you (OP) provided
scatterlike(A, I)
array([[[ 1,  2,  0],
        [ 4,  5,  6],
        [ 7,  8,  9]],

       [[10, 11, 12],
        [13, 14, 15],
        [16,  0, 18]],

       [[19, 20, 21],
        [22,  0, 24],
        [25, 26, 27]]])

# Higher dimensions, B.shape is (2,3,3,3)
scatterlike(B, C)
array([[[[ 1,  2,  0],
         [ 4,  5,  6],
         [ 7,  8,  9]],

        [[10, 11, 12],
         [13, 14, 15],
         [16,  0, 18]],

        [[19, 20, 21],
         [22,  0, 24],
         [25, 26, 27]]],


       [[[ 0,  2,  3],
         [ 4,  5,  6],
         [ 7,  8,  9]],

        [[10, 11, 12],
         [13,  0, 15],
         [16, 17, 18]],

        [[19, 20, 21],
         [22, 23, 24],
         [25, 26,  0]]]])

# Even higher dimensions, B[None].repeat(2,0) is of shape (2,2,3,3,3)
# Output should be like above but repeated
scatterlike(B[None].repeat(2,0), C[None].repeat(2,0))
array([[[[[ 1,  2,  0],
          [ 4,  5,  6],
          [ 7,  8,  9]],

         [[10, 11, 12],
          [13, 14, 15],
          [16,  0, 18]],

         [[19, 20, 21],
          [22,  0, 24],
          [25, 26, 27]]],


        [[[ 0,  2,  3],
          [ 4,  5,  6],
          [ 7,  8,  9]],

         [[10, 11, 12],
          [13,  0, 15],
          [16, 17, 18]],

         [[19, 20, 21],
          [22, 23, 24],
          [25, 26,  0]]]],



       [[[[ 1,  2,  0],
          [ 4,  5,  6],
          [ 7,  8,  9]],

         [[10, 11, 12],
          [13, 14, 15],
          [16,  0, 18]],

         [[19, 20, 21],
          [22,  0, 24],
          [25, 26, 27]]],


        [[[ 0,  2,  3],
          [ 4,  5,  6],
          [ 7,  8,  9]],

         [[10, 11, 12],
          [13,  0, 15],
          [16, 17, 18]],

         [[19, 20, 21],
          [22, 23, 24],
          [25, 26,  0]]]]])

奖金

只是为了好玩，我认为您甚至可以将该函数用于 (*S,*G) (通过从 I 获取形状信息) ，也就是说，您可以使用 (N x M x L x ...) 来代替形状为 (N x N) 的二维数组。您还可以指定每个目标值应该是什么:

from numpy.typing import ArrayLike
from typing import Union
def scatterlike_general(A: np.ndarray, I: np.ndarray, target: Union[ArrayLike,float] = 0):
    A_ = A.reshape(-1, *A.shape[-I.shape[-1]:])
    A_[(range(len(A_)),*I.reshape(-1,I.shape[-1]).T.tolist())] = target
    return A_.reshape(A.shape)

需要注意的是，我还没有测试过scatterlike_general，但我认为它应该可以工作。