如何利用Numpy高效按条件分割数组（提取非零连续块）

嘿，这个需求用Numpy的向量化操作来实现绝对比for循环简洁高效多了，尤其是处理大型数组的时候，性能提升会非常明显。我给你写个实现，顺便拆解下逻辑：

首先，核心思路是利用Numpy的索引和数组分割功能，跳过循环直接定位并提取连续块：

import numpy as np

def chunks(array):
    # 第一步：获取所有非零元素的索引
    non_zero_idx = np.where(array != 0)[0]
    
    # 处理数组全为0的边界情况
    if len(non_zero_idx) == 0:
        return []
    
    # 第二步：找到连续块之间的分割点
    # 相邻索引差大于1的位置，就是两个块的分界处
    split_positions = np.where(np.diff(non_zero_idx) > 1)[0] + 1
    
    # 第三步：把索引数组分割成对应每个连续块的子数组
    index_chunks = np.split(non_zero_idx, split_positions)
    
    # 第四步：生成要求的结果格式：[起始索引, 块内容列表]
    result = [[chunk[0], array[chunk].tolist()] for chunk in index_chunks]
    
    return result

# 测试你的示例输入
test_array = np.array([1, 0, 0, 0, 1, 1, 0, 0, 1, 0])
print(chunks(test_array))
# 输出：[[0, [1]], [4, [1, 1]], [8, [1]]]

代码逻辑拆解： #

• np.where(array != 0)[0]：快速定位所有非零元素的索引，这一步是向量化操作，比Python循环遍历数组快几个数量级。
• np.diff(non_zero_idx)：计算相邻非零索引的差值，连续的非零元素索引差为1，差值大于1的地方就代表两个连续块之间的间隔。
• np.split(non_zero_idx, split_positions)：把索引数组按分割点拆分成多个子数组，每个子数组对应一个连续非零块的索引集合。
• 最后用列表推导式把每个块的起始索引（子数组第一个元素）和对应的元素列表组合起来，tolist()把Numpy数组转换成普通Python列表，完全匹配你想要的输出格式。

这种写法不仅代码更简洁，而且完全利用了Numpy的底层优化，处理大数组时的性能碾压for循环版本～

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