NumPy如何按批量起止索引从二维通道数组提取片段生成三维数组

输入输出规则说明 #

• 输入data：形状为[chn, dim]的二维NumPy数组，chn代表通道总数，每个通道对应存储长度为dim的一维连续数据
• 输入indices：形状为[seg, 2]的二维NumPy数组，每一行对应一个待提取片段的索引边界；其中结束索引为闭区间（提取时会覆盖该位置的数据），例如片段[3,5]代表提取索引3、4、5三个位置的数据
• 输出extracted：形状为[seg, chn, 单段长度]的三维NumPy数组，三个维度依次为提取得到的片段总数、数据通道数、每个片段的连续提取长度。

函数实现 #

利用NumPy的广播索引与轴调整机制实现向量化提取，避免逐元素循环的效率损耗：

import numpy as np

def do_something(data: np.ndarray, indices: np.ndarray) -> np.ndarray:
    # 拆分所有片段的起始、结束边界
    starts = indices[:, 0]
    ends = indices[:, 1]
    # 按闭区间规则计算单段提取长度（题目要求输出为定长三维数组，默认所有片段长度一致）
    seg_len = ends[0] - starts[0] + 1
    # 构造每个片段在数据维度的取数位置索引，形状为 [seg, seg_len]
    dim_take_pos = starts[:, None] + np.arange(seg_len)
    # 索引取数得到初始形状 [chn, seg, seg_len]，再调整轴顺序为目标形状 [seg, chn, seg_len]
    extracted = data[:, dim_take_pos].transpose(1, 0, 2)
    return extracted

测试验证 #

使用题目给出的示例代码运行，输出和预期完全一致：

# 构造测试数据
data1 = np.expand_dims(np.arange(0, 100), 0)
data2 = np.expand_dims(np.arange(100, 200), 0)
data3 = np.expand_dims(np.arange(200, 300), 0)
data = np.concatenate([data1, data2, data3])
indices = np.array([[3,5], [8,10], [13,15], [16,18]])

extracted = do_something(data, indices)
print("extracted.shape", extracted.shape) # 输出：extracted.shape (4, 3, 3)
print(extracted)

运行后的打印结果与题目预期完全匹配：

[[[  3   4   5]
  [103 104 105]
  [203 204 205]]

 [[  8   9  10]
  [108 109 110]
  [208 209 210]]

 [[ 13  14  15]
  [113 114 115]
  [213 214 215]]

 [[ 16  17  18]
  [116 117 118]
  [216 217 218]]]

实现说明 #

• 索引逻辑严格遵循闭区间规则，单段长度计算为end - start + 1，和题目示例规则完全匹配
• 全程采用NumPy原生向量化操作，相比逐片段循环提取的写法效率更高，尤其适用于通道数、片段数较多的场景
• 输出数组的维度顺序严格符合要求：第一维为片段总数、第二维为通道数、第三维为单段提取长度

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