时间戳

首页 > 解决方案 > Numpy 中的类散点操作

python - Numpy 中的类散点操作

问题描述

我面临以下问题：

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

我想设置

A[i1, ...., ik, I[i1, ..., ik, 0], I[i1, ...., ik, 1]] := 0（伪代码）

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

例如，考虑以下示例 whereN=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]]])

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

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

我将不胜感激建议！

标签： pythonnumpy

解决方案

有趣的问题！

我相信您可以简单地分别重塑A和I到(-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 N)，您可以拥有 like ，而不是只有 shape 的 2D 数组(N x M x L x ...)。您还可以指定每个目标值应该是什么：

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，但我认为它应该可以工作。

推荐阅读