高效填充 NumPy 数组中的零值位置：向量化实现无循环批量赋值

本文介绍一种纯 numpy 向量化方法，将源数组 y 中每行的非零元素，按顺序填入目标数组 x 每行最靠前的零值位置，全程避免 python 循环，适用于大规模矩阵（如 5000×20000）高频操作。

本文介绍一种纯 numpy 向量化方法，将源数组 y 中每行的非零元素，按顺序填入目标数组 x 每行最靠前的零值位置，全程避免 python 循环，适用于大规模矩阵（如 5000×20000）高频操作。

在科学计算与数据预处理中，常需将一个数组的“有效值”精准注入另一个数组的空白槽位（即零值位置），且要求严格保持行内顺序与位置优先级。典型场景包括稀疏特征对齐、掩码填充、增量更新等。当矩阵规模达万级行列且操作频次极高时，Python 显式循环会成为严重性能瓶颈。本文提供一种完全向量化、无 for 循环、内存友好的 NumPy 实现方案。

核心思路：布尔索引 + 累计计数定位

关键在于两点：

Replit Ghostwrite

一种基于 ML 的工具，可提供代码完成、生成、转换和编辑器内搜索功能。

• 识别“空位”：X == 0 得到布尔掩码 i1，标记所有可写入位置；
• 定位“前 kᵢ 个空位”：利用 i1.cumsum(1) 对每行从左到右累计零值个数，再与每行非零元数量 i2.sum(1, keepdims=True) 比较，构造精确写入掩码。

具体逻辑如下：

1. i1 = X == 0：获取 X 中所有零值位置；
2. i2 = Y != 0：获取 Y 中所有非零值位置（即待复制的元素）；
3. i2.sum(1, keepdims=True)：计算每行 Y 的非零元素个数 k_i，保持列维度；
4. i1.cumsum(1)
5. 最终掩码 (i1.cumsum(1)
6. Y[i2] 按行优先顺序拉平所有非零元素，恰好与目标位置一一对应（NumPy 高级索引的扁平化规则保障此一致性）。

完整实现与示例

import numpy as np

def fill_empty_np(X: np.ndarray, Y: np.ndarray) -> None:
    """
    将 Y 中每行的非零元素，依次填入 X 中对应行最左侧的零值位置。

    要求：X.shape == Y.shape，且对每行 i，X[i] 中零的个数 >= Y[i] 中非零个数。
    原地修改 X，不返回新数组。
    """
    i1 = X == 0
    i2 = Y != 0
    # 构造写入掩码：仅在 X 的零位中，且是前 k_i 个零位的位置赋值
    mask = (i1.cumsum(axis=1) <= i2.sum(axis=1, keepdims=True)) & i1
    X[mask] = Y[i2]  # Y[i2] 自动展平，与 mask 中 True 位置按行优先顺序对齐

# 示例验证
X = np.array([[1, 0, 1, 1, 0],
              [1, 0, 0, 1, 1],
              [1, 0, 0, 0, 0]])

Y = np.array([[0, 0, 2, 0, 0],
              [0, 0, 0, 2, 2],
              [2, 0, 2, 2, 0]])

fill_empty_np(X, Y)
print(X)
# 输出：
# [[1 2 1 1 0]
#  [1 2 2 1 1]
#  [1 2 2 2 0]]

注意事项与性能提示

• ✅ 严格形状匹配：X 与 Y 必须同形，否则 cumsum 和 sum(..., keepdims=True) 广播失败；
• ✅ 容量保证：算法不校验 X 每行零位是否足够 —— 若 m_i
• ⚠️ 内存开销：生成多个临时布尔/整型数组（i1, i2, cumsum, mask），对超大矩阵（如 5000×20000）可能占用数 GB 内存；若内存受限，建议分块处理或改用 Numba（如答案中提及）；
• ⚡ 性能对比：该方案比纯 Python 循环快 10–100 倍，但比 Numba JIT 编译版本慢约 2–5 倍（因未短路、全量计算）；如追求极致性能且可引入额外依赖，推荐搭配 @njit(parallel=True) 重写；
• ? 不可逆操作：函数原地修改 X；如需保留原始 X，调用前使用 X_copy = X.copy()。

掌握此模式，你不仅能高效解决当前问题，更能迁移应用于各类“按条件定位+批量赋值”的 NumPy 工程场景。