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

Changing cells in a 4D numpy matrix subject to conditions on axis indexes

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

如何将 A 中满足某些条件的所有单元格设置为 0?有没有一种方法可以在没有循环的情况下有效地做到这一点(使用矢量化?)。

示例条件:具有第 3 维索引 k 的所有单元格,其中 B[k] == x

例如,

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

然后对于 A 的第二个维度(即列),我想将第二个维度中的所有单元格归零,从而使该维度中的单元格索引(调用索引 i ), 满足条件 B[i] == 7。在这种情况下,A 将转换为 A = [[0,0],[3,4]].

以下有帮助吗?

A = np.arange(16,dtype='float64').reshape(2,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, j, k, l))
b_k = np.random.random(k)
b_j = np.random.random(j)

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

# You can alsow combine multiple conditions along the different axes
# i,         j,         k, l
a[:, b_j > 0.5, b_k < 0.5, :] = 0

# Or work with the index explicitly
condition_k = np.arange(k) % 3 == 0  # Is the index divisible by 3?
# i, j,            k, l
a[:, :,  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]])