沿数组中的轴传播真实条目

Question

我必须多次执行以下操作。使用 numpy 函数而不是循环，我通常会获得非常好的性能，但我无法将其复制到更高维的数组。任何建议或替代方案将非常受欢迎：

我有一个布尔数组，我想将 true indeces 传播到接下来的 2 个位置，例如：

如果这个一维数组（A）是：

import numpy as np

# Number of positions to propagate the array
propagation = 2

# Original array
A = np.array([False, True, False, False, False, True, False, False, False, False, False, True, False])

我可以创建一个“空”数组，然后找到索引，传播它们，然后展平 argwhere，然后展平它：

B = np.zeros(A.shape, dtype=bool)

# Compute the indeces of the True values and make the two next to it True as well
idcs_true = np.argwhere(A) + np.arange(propagation + 1)
idcs_true = idcs_true.flatten()
idcs_true = idcs_true[idcs_true < A.size] # in case the error propagation gives a great
B[idcs_true] = True

# Array
print(f'Original array     A = {A}')
print(f'New array (2 true) B = {B}')

给出：

Original array     A = [False  True False False False  True False False False False False  True
 False]
New array (2 true) B = [False  True  True  True False  True  True  True False False False  True
  True]

但是，这会变得更加复杂并且会失败，例如：

AA = np.array([[False, True, False, False, False, True, False, False, False, False, False, True, False],
               [False, True, False, False, False, True, False, False, False, False, False, True, False]])

感谢您的任何建议。

Answer 1

您可以将当前代码包装在一个函数中，然后将其与二维 AA 数组一起提供给

np.apply_along_axis

函数。

import numpy as np

# Number of positions to propagate the array
propagation = 2

# Original array
AA = np.array([[False, True, False, False, False, True, False, False, False, False, False, True, False],
             [False, True, False, False, False, True, False, False, False, False, False, True, False]])

def prop_func(A):
    B = np.zeros(A.shape, dtype=bool)

    # Compute the indices of the True values and make the two next to it True as well
    idcs_true = np.argwhere(A) + np.arange(propagation + 1)
    idcs_true = idcs_true.flatten()
    idcs_true = idcs_true[idcs_true < A.size] # in case the error propagation gives a great
    B[idcs_true] = True
    return B

# Apply prop_func along the rows (axis=1)
BB = np.apply_along_axis(prop_func, 1, AA)

# Array
print(f'Original array     AA = {AA}')
print(f'New array (2 true) BB = {BB}')

Original array     AA = 
[[False  True False False False  True False False False False False  True
  False]
 [False  True False False False  True False False False False False  True
  False]]
New array (2 true) BB = 
[[False  True  True  True False  True  True  True False False False  True
   True]
 [False  True  True  True False  True  True  True False False False  True
   True]]