NumPy中如何按条件替换ndarray元素并调整行内元素位置

我有如下示例数组：

[[ 2  4 -1]
 [ 1 -1 -1]
 [ 2  4 -1]
 [ 0 -1 -1]
 [ 0  0  2]
 [ 2  4 -1]]

需求为：针对每一行，若该行存在两个0，则将第二个0替换为-1，并将其移动到该行末尾，处理后的预期效果如下：

[[ 2  4 -1]
 [ 1 -1 -1]
 [ 2  4 -1]
 [ 0 -1 -1]
 [ 0  2 -1]
 [ 2  4 -1]]

我尝试编写了如下循环代码实现该需求：

for i in np.nditer(f_vars_r, op_flags = ['readwrite']):
    if i == 0 and i==0:
        i[...] == -1

我清楚这段代码逻辑并不完整，但至少预期会触发报错，可奇怪的是代码运行无任何报错，也没有对数组做出修改。

• 条件判断i == 0 and i==0是重复判断，仅校验当前遍历的单个元素是否为0，既无法统计单行内0的总数，也定位不到第二个0的位置
• 语句i[...] == -1用了双等号，是比较运算而非赋值操作，只会返回布尔值False，不会修改原数组元素
• np.nditer默认逐元素遍历数组，没有按行聚合的能力，本身就不适合实现行级规则处理

逐行循环实现（逻辑直观易理解） #

直接遍历数组每一行处理，不需要用np.nditer：

import numpy as np

# 示例数组初始化
f_vars_r = np.array([
    [2,4,-1],
    [1,-1,-1],
    [2,4,-1],
    [0,-1,-1],
    [0,0,2],
    [2,4,-1]
])

for row in f_vars_r:
    # 找到当前行所有0的索引
    zero_positions = np.where(row == 0)[0]
    # 仅当行内至少2个0时处理
    if len(zero_positions) >= 2:
        second_zero_idx = zero_positions[1]
        # 删除第二个0，末尾追加-1
        row[:] = np.concatenate([np.delete(row, second_zero_idx), [-1]])

向量化筛选实现（大数组性能更优） #

先批量筛选出需要处理的行，再做修改，比逐行全量遍历效率更高：

# 统计每行0的数量
zero_counts = (f_vars_r == 0).sum(axis=1)
# 定位到需要处理的行索引
need_process_rows = np.where(zero_counts >= 2)[0]

for idx in need_process_rows:
    current_row = f_vars_r[idx]
    second_zero = np.where(current_row == 0)[0][1]
    f_vars_r[idx] = np.concatenate([np.delete(current_row, second_zero), [-1]])

运行上述任意一段代码后，打印数组即可得到预期输出：

[[ 2  4 -1]
 [ 1 -1 -1]
 [ 2  4 -1]
 [ 0 -1 -1]
 [ 0  2 -1]
 [ 2  4 -1]]

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