Numpy 中的 Scatter-Like 操作

如何解决Numpy 中的 Scatter-Like 操作

我面临以下问题：

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

我想设置

A[i1,....,ik,I[i1,0],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,18]]])

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

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

我将不胜感激！

解决方法

有趣的问题！

我相信您可以简单地将 A 和 I 分别重塑为 (-1,N,N) 和 (-1,2)，并始终处理 (K x N x N) 中的问题，然后只需 {{1 }} 回到原来的维度。我提出的解决方案如下：

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,1]],[[0,0],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,M) 类型建立索引，因为 np.ndarray 和其他可迭代对象的索引行为不同。

以下是一些输出：

np.ndarray

奖金

只是为了好玩，我认为您甚至可以将函数用于 # Simplest case scatterlike(A[0],I[0]) array([[1,[4,[7,9]]) # Extended version of the example you (OP) provided scatterlike(A,I) array([[[ 1,27]]]) # Higher dimensions,B.shape is (2,3,3) scatterlike(B,C) array([[[[ 1,27]]],[[[ 0,0]]]]) # Even higher dimensions,B[None].repeat(2,0) is of shape (2,3) # Output should be like above but repeated scatterlike(B[None].repeat(2,0),C[None].repeat(2,0)) array([[[[[ 1,0]]]],[[[[ 1,0]]]]]) （通过从 (*S,*G) 获取形状信息），也就是说，不仅仅是形状 {{}} 的二维数组{1}}，您可以像 I。您也可以指定每个目标值应该是什么：

(N x N)

需要注意的是，我没有测试过 (N x M x L x ...)，但我认为它应该可以工作。

,

您可以使用扁平网格访问由 S 表示的 A 的尺寸。 要访问 I 中的适当位置，您可以使用相同的展平网格。 他是一些实现这个想法的代码。

这是一个额外的更复杂的例子（抱歉没有在输出中包含整个数组 A）：