Numpy - 在多个轴上设置第一个最大值

Question

我有一个多维数组，最后两个维度对应一个“窗口”。对于数组中的每个“窗口”，我的目标是将窗口中的第一个最大值设置为 1，将其余值设置为 0。例如，对于一个窗口，我们有以下赋值：

[[0., 0., 0.],          [[0., 0., 0.],
 [0., 1., 1.],    ->     [0., 1., 0.],
 [0., 1., 1.]],          [0., 0., 0.]]

现在，我想对所有窗口执行此操作，其中每个窗口都在最后两个维度上，即axis=(-1,-2)。这是我迄今为止尝试过的：

# windows
x = np.array([[[[[[0., 0., 0.],
                  [0., 1., 1.],
                  [0., 1., 1.]],

                 [[0., 0., 0.],
                  [1., 1., 0.],
                  [1., 1., 0.]]],

  
                [[[0., 1., 1.],
                  [0., 1., 1.],
                  [0., 1., 1.]],

                 [[1., 1., 0.],
                  [1., 1., 0.],
                  [1., 1., 0.]]],


                [[[0., 1., 1.],
                  [0., 1., 1.],
                  [0., 1., 0.]],

                 [[1., 1., 0.],
                  [1., 1., 0.],
                  [1., 0., 1.]]]]]])

max_indices = np.argwhere(x == x.max()) # all indicies of max values in x
mask = np.zeros_like(x) # start of with everything set to 0

# I can set the first window by doing the following
mask[tuple(max_indices[0])] = 1.0

我怎样才能对所有其他窗口做同样的事情？最好不要循环。

我的想法是在不考虑最后两个维度的 max_indices 中找到唯一值，但我不知道该怎么做。

编辑

我有一个案例，上面场景中的x 取自另一个 数组使用np.lib.stride_tricks.as_strided，这意味着步幅会有所不同。 情况如下：

# The image or whatever i want to modify
x = np.array([[[[0., 0., 0., 0.],
                [0., 1., 1., 0.],
                [0., 1., 1., 0.],
                [0., 1., 1., 0.],
                [0., 1., 0., 1.]]]])

# the windows of that image
x = np.lib.stride_tricks.as_strided(x, shape=(1, 1, 3, 2, 3, 3), strides=(160, 160, 32, 8, 32, 8))

使用Felipe给出的解决方案：

y = x[0][0]  # Just unnesting it a bit
yr = y.reshape(-1, 9)
idx = yr.argmax(1)
y0 = np.zeros_like(yr)
np.put_along_axis(y0, idx[:, None], 1, axis=1)
y = y0.reshape(y.shape)

# plug it back
x[0][0] = y

x

array([[[[[[0., 0., 0.],
           [0., 1., 0.],
           [0., 1., 0.]],

          [[0., 0., 0.],
           [1., 0., 0.],
           [1., 0., 0.]]],


         [[[0., 1., 0.],
           [0., 1., 0.],
           [0., 0., 0.]],

          [[1., 0., 0.],
           [1., 0., 0.],
           [0., 0., 0.]]],


         [[[0., 1., 0.],
           [0., 0., 0.],
           [0., 0., 0.]],

          [[1., 0., 0.],
           [0., 0., 0.],
           [0., 0., 0.]]]]]])

形状相同，但在x 和y 中的步幅不同，似乎“插回去”会导致 x 和 y 不相同，这是我所预料的。

Answer 1

这是使用np.put_along_axis 的一种方法（尽管这实际上只是一个更快的循环）。

y = x[0][0]  # Just unnesting it a bit
yr = y.reshape(-1, 9)
idx = yr.argmax(1)
y0 = np.zeros_like(yr)
np.put_along_axis(y0, idx[:, None], 1, axis=1)
y = y0.reshape(y.shape)

输出：

array([[[[0., 0., 0.],
         [0., 1., 0.],
         [0., 0., 0.]],

        [[0., 0., 0.],
         [1., 0., 0.],
         [0., 0., 0.]]],


       [[[0., 1., 0.],
         [0., 0., 0.],
         [0., 0., 0.]],

        [[1., 0., 0.],
         [0., 0., 0.],
         [0., 0., 0.]]],


       [[[0., 1., 0.],
         [0., 0., 0.],
         [0., 0., 0.]],

        [[1., 0., 0.],
         [0., 0., 0.],
         [0., 0., 0.]]]])

编辑

不是最聪明的想法，但我的意思是你可以先将它重新分配给它的副本，然后再次获得你的跨步视图，例如：

x0 = x.copy()  # Get a normal view

# [ Run the previous snippet ]

x0[0][0] = y  # Assign the modified version

# Then if you really want another strided view you can recreate it
x0 = x0.reshape(saved_x.shape)  # Reshape it to the original version pre-stride
x = np.lib.stride_tricks.as_strided(x0, shape=(1, 1, 3, 2, 3, 3), strides=(160, 160, 32, 8, 32, 8))