基于条件修改3D NumPy数组值：解决索引定位错误问题

问题分析 #

你需要在每个3D块中，找到包含两个0的列，并将该列中行索引较小的0替换为-1。原代码仅针对固定行（第2、3行）的0进行判断，若需求是任意位置存在两个0的列，则逻辑不够通用；若仅针对固定行的情况，原代码其实可以得到正确结果，以下提供两种解决方案适配不同场景。

方案1：针对固定行（第2、3行）的简化解法 #

如果你的需求明确是每个块的第2行（索引1）和第3行（索引2）同时为0的列，直接用以下代码即可：

import numpy as np

# 原始数组
a = np.array([[[-2, -1],
               [-1,  0],
               [ 0,  0]],
              [[-1, -1],
               [-1,  0],
               [ 0,  0]]])

# 找出每个块中，第2行和第3行都为0的列
mask = (a[:, 1, :] == 0) & (a[:, 2, :] == 0)

# 将对应块的第2行、符合条件的列设为-1
a[:, 1, mask] = -1

print(a)
# 输出与期望一致：
# [[[-2 -1]
#   [-1 -1]
#   [ 0  0]]
#  [[-1 -1]
#   [-1 -1]
#   [ 0  0]]]

方案2：通用解法（任意位置含两个0的列） #

如果需要处理任意行中存在两个0的列，并自动找到行索引最小的0替换，可通过向量化操作实现：

import numpy as np

a = np.array([[[-2, -1],
               [-1,  0],
               [ 0,  0]],
              [[-1,  0],
               [ 0,  0],
               [-1, -1]]])

# 1. 标记所有0的位置
zero_mask = (a == 0)

# 2. 计算每个块每一列的0的数量
zero_counts = zero_mask.sum(axis=1)

# 3. 筛选出含至少两个0的列
cols_to_modify = zero_counts >= 2

# 4. 找出每个块每一列中第一个0的行索引（行索引最小的0）
first_zero_row = zero_mask.argmax(axis=1)

# 5. 获取需要修改的位置索引
block_idx, col_idx = np.where(cols_to_modify)
row_idx = first_zero_row[block_idx, col_idx]

# 6. 执行替换
a[block_idx, row_idx, col_idx] = -1

print(a)
# 输出示例（第二个块的第一列有两个0，替换第一个0）：
# [[[-2 -1]
#   [-1 -1]
#   [ 0  0]]
#  [[-1 -1]
#   [ 0  0]
#   [-1 -1]]]

内容的提问来源于stack exchange，提问作者jasmine