根据轴索引的条件更改4D numpy矩阵中的像元

问题描述

假设我有一个4D numpy数组A，其中四个维度的索引为i,j,k,l，假设50 x 40 x 30 x 20。还要假设我还有其他列表B。

如何将A中满足某些条件的所有单元格设置为0？有没有一种方法可以有效地做到无循环（使用矢量化？）。

示例条件：所有具有第三维索引k的单元格，其中B[k] == x

例如，

如果我们有2D矩阵A = [[1,2],[3,4]]和B = [7,8]

然后对于A的第二维（即列），我想将第二维中的所有单元格归零，从而该单元格在该维中的索引（称为索引i），满足条件B[i] == 7。在这种情况下，A将转换为 A = [[0,0],4]]。

解决方法

以下帮助吗？

A = np.arange(16,dtype='float64').reshape(2,2,2)
A[A == 2] = 3.14

我用3.14替换了等于2的条目。您可以将其设置为其他值。

,

您可以为特定轴指定布尔数组：

import numpy as np

i,j,k,l = 50,40,30,20
a = np.random.random((i,l))
b_k = np.random.random(k)
b_j = np.random.random(j)

# i,l
a[:,:,b_k < 0.5,:] = 0

# You can alsow combine multiple conditions along the different axes
# i,b_j > 0.5,:] = 0

# Or work with the index explicitly
condition_k = np.arange(k) % 3 == 0  # Is the index divisible by 3?
# i,condition_k,:] = 0

要使用您给出的示例

a = np.array([[1,2],[3,4]])
b = np.array([7,8])

#      i,j
a[b == 7,:] = 0
# array([[0,0],#        [3,4]])