如何基于负值条件分割NumPy数组并移除负值

针对大规模数组的场景，我们可以利用NumPy的向量化操作和索引功能，高效完成“移除负值+按原负值位置分割”的需求，具体实现如下：

import numpy as np

arr = np.array([1, 2, 3, 4, -5, -6, 3, 5, 1, -2, 5, -1, -1, 10])

# 标记所有负值的位置
neg_mask = arr < 0
# 获取负值的索引，并生成分割边界（包含数组首尾）
split_boundaries = np.concatenate(([0], np.where(neg_mask)[0] + 1, [len(arr)]))

# 提取非负值数组及其在原数组中的索引
non_neg_vals = arr[~neg_mask]
non_neg_indices = np.where(~neg_mask)[0]

# 找到每个分割边界在非负值数组中的对应位置
start_indices = np.searchsorted(non_neg_indices, split_boundaries[:-1], side='left')
end_indices = np.searchsorted(non_neg_indices, split_boundaries[1:], side='left')

# 分割并过滤空列表
split_list = [non_neg_vals[s:e].tolist() for s, e in zip(start_indices, end_indices)]
split_list = [lst for lst in split_list if lst]

print(split_list)
# 输出：[[1, 2, 3, 4], [3, 5, 1], [5], [10]]

关键步骤说明 #

• 标记与边界生成：通过neg_mask定位负值，将数组起始点、每个负值的下一个索引（分割点）、数组终点拼接成完整的分割边界。
• 非负值映射：提取非负值数组和它们在原数组中的位置，为后续匹配分割边界做准备。
• 快速索引匹配：用np.searchsorted高效找到分割边界在非负值数组中的对应索引，这一步是实现向量化、提升效率的核心，避免了Python层面的循环。
• 分割与过滤：根据索引分割非负值数组，转成列表后过滤掉连续负值导致的空区间。

该方案全程依赖NumPy的底层优化操作，相比列表推导式，在处理大规模数组时速度提升显著，适合重复执行的计算场景。

内容的提问来源于stack exchange，提问作者Alex V.