NumPy 中实现不规则嵌套数组的逐元素减法：高效向量化方案-Python教程-PHP中文网

NumPy 中实现不规则嵌套数组的逐元素减法：高效向量化方案

本文介绍如何在 NumPy 中对形状不一致的嵌套数组（如 (N, 2) 与 [(M₁,2), (M₂,2), ...]）进行高效、无显式循环的逐元素减法，核心是利用 np.concatenate + np.repeat + np.split 完成结构对齐与向量化计算。

本文介绍如何在 numpy 中对形状不一致的嵌套数组（如 `(n, 2)` 与 `[(m₁,2), (m₂,2), ...]`）进行高效、无显式循环的逐元素减法，核心是利用 `np.concatenate` + `np.repeat` + `np.split` 完成结构对齐与向量化计算。

在 NumPy 中，标准广播机制无法直接处理“一个二维数组与一个长度为 N 的对象数组（每个元素是形状为 (Mᵢ, 2) 的二维数组）”之间的逐行减法。例如，给定：

a = np.array([[1, 2], [3, 4]]) —— 形状 (2, 2)
b_list = [np.array([[1, 2]]), np.array([[1, 2], [3, 4]])] —— 各子数组形状分别为 (1, 2) 和 (2, 2)

期望结果是 [a[0] - b_list[0], a[1] - b_list[1]]，即：

[array([[0, 0]]), array([[2, 2], [0, 0]])]

但直接写 a - b_list 会触发对象数组的逐元素广播（生成 (2, 2) 个对象），导致语义错误；而 np.subtract(a, b_list) 同样不支持 axis 参数，也无法适配变长结构。

✅ 正确解法是结构对齐 + 向量化运算 + 结构还原三步策略：

拼接（Concatenate）：将所有 b_list[i] 沿第 0 轴堆叠为单一大数组 b_concat；
重复（Repeat）：将 a[i] 按 len(b_list[i]) 次数沿第 0 轴重复，得到 a_repeat，使其与 b_concat 长度一致；
分割（Split）：对向量化差值 diff_flat = a_repeat - b_concat，按原始各段长度切分回列表。

完整可运行示例：

import numpy as np

a = np.array([[1, 2], [3, 4]])
b_list = [np.array([[1, 2]]), np.array([[1, 2], [3, 4]])]

# Step 1: 拼接 b_list → (3, 2)
b_concat = np.concatenate(b_list, axis=0)

# Step 2: 重复 a[i] 对应次数 → (3, 2)
repeats = [len(b) for b in b_list]  # [1, 2]
a_repeat = np.repeat(a, repeats, axis=0)

# Step 3: 向量化减法 & 分割还原
diff_flat = a_repeat - b_concat
split_indices = np.cumsum(repeats)[:-1]  # [1]
diff_final = np.split(diff_flat, split_indices)

print(diff_final)
# [array([[0, 0]]), array([[2, 2], [0, 0]])]

⚠️ 注意事项：

所有 b_list[i] 必须具有相同最后一维（如均为 2 列），否则 concatenate 会报错；
np.repeat(..., axis=0) 要求 a 是常规 ndarray（非 object dtype），确保类型安全；
若 b_list 来自不规则数据（如不同长度的轨迹点），建议预先校验维度一致性；
该方法显著优于列表推导（实测提速约 25%），尤其在 N 较大时（如 N=300, M=30 场景下，500 次调用耗时从 3.34s 降至 2.54s）。