如何仅使用NumPy函数将指定格式一维np.array拆分为三个对应子序列

可以仅用NumPy实现该需求，不需要使用for循环。核心逻辑是定位数组中第一个和最后一个非零元素的索引，直接按索引切片即可得到三个子序列。

代码实现 #

import numpy as np

def split_array(arr: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    # 获取所有非零元素的索引数组
    non_zero_idx = np.nonzero(arr)[0]
    # 取第一个、最后一个非零元素的索引
    start, end = non_zero_idx[0], non_zero_idx[-1]
    # 切片得到三个子序列
    seq1 = arr[:start]
    seq2 = arr[start:end+1]
    seq3 = arr[end+1:]
    return seq1, seq2, seq3

测试验证 #

用题目给出的三个示例验证效果：

# 示例1
arr1 = np.array([0,0,0,1,1,1,1,0,0,0])
seq1, seq2, seq3 = split_array(arr1)
print(seq1) # 输出 [0 0 0]
print(seq2) # 输出 [1 1 1 1]
print(seq3) # 输出 [0 0 0]

# 示例2
arr2 = np.array([0,1,0,0,0,0,0,0,1,0])
seq1, seq2, seq3 = split_array(arr2)
print(seq1) # 输出 [0]
print(seq2) # 输出 [1 0 0 0 0 0 0 1]
print(seq3) # 输出 [0]

# 示例3
arr3 = np.array([0,1,1,1,0,1,0,1,1,0])
seq1, seq2, seq3 = split_array(arr3)
print(seq1) # 输出 [0]
print(seq2) # 输出 [1 1 1 0 1 0 1 1]
print(seq3) # 输出 [0]

逻辑说明 #

• 题目规则已经保证数组中至少存在一个1，因此np.nonzero返回的索引数组不会为空，无需额外处理边界异常
• np.nonzero是NumPy原生函数，时间复杂度为O(n)，效率远高于Python层面的手动遍历
• 索引切片也是NumPy原生操作，全程没有引入任何for循环

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