如何将NumPy数组拆分为相同元素组成的子数组？

这里提供两种实用的实现方式，满足你的需求：

方法一：纯NumPy原生实现（适合处理大型数组） #

利用NumPy的数组操作定位分割点，效率更高：

import numpy as np

def stackoverflowgeniusfunc(arr):
    # 定位连续元素发生变化的位置
    split_indices = np.where(arr[1:] != arr[:-1])[0] + 1
    # 拆分原数组
    sub_arrays = np.split(arr, split_indices)
    # 转换为列表格式（若保留NumPy子数组可去掉此步）
    return [sub.tolist() for sub in sub_arrays]

# 测试示例
A = np.array([1,2,2,2,3,3,1,1,1,4,4,4,4,4,5,5])
B = stackoverflowgeniusfunc(A)
print(B)
# 输出：[[1], [2, 2, 2], [3, 3], [1, 1, 1], [4, 4, 4, 4, 4], [5, 5]]

方法二：使用itertools.groupby（代码更简洁） #

借助Python标准库的groupby工具，快速对连续相同元素分组：

from itertools import groupby
import numpy as np

def stackoverflowgeniusfunc(arr):
    # groupby会自动聚合连续相同元素，直接转为列表即可
    return [list(group) for _, group in groupby(arr)]

# 测试示例
A = np.array([1,2,2,2,3,3,1,1,1,4,4,4,4,4,5,5])
B = stackoverflowgeniusfunc(A)
print(B)
# 输出：[[1], [2, 2, 2], [3, 3], [1, 1, 1], [4, 4, 4, 4, 4], [5, 5]]

两种方法对比 #

• 纯NumPy方法：更适合处理大规模数组，计算效率优于迭代器方式
• groupby方法：代码更简洁，无需复杂的数组操作，同时支持普通Python列表输入

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