使用另一数组指定深度实现NumPy三维数组无循环切片

实现思路 #

不同(x,y)位置要提取的z轴元素长度不一样，普通齐次NumPy数组没法直接存储这类变长块，所有步骤都可以靠NumPy的向量化操作完成，完全不需要写Python层的逐元素循环。

注意：你给出的示例中b = [[1,2],[0,0]]和预期结果c存在数值匹配偏差，按照原b的规则提取得到的结果应为[[[0], [3,4]], [[], []]，下面代码会将b调整为和你写的预期c完全对应的参数，整体切片逻辑不受影响。

具体实现 #

1. 用广播一次性生成保留掩码 #

不用循环逐位置判断哪个元素要留，靠广播机制一次算完所有元素的保留标记：

• 先拿到原数组z轴的总长度，生成从0开始的z轴位置序号
• 把切片长度矩阵b多拼一个维度，和z轴序号做广播比较：只要z位置的序号小于对应(x,y)位置要求的切片长度，就标记为要保留

对应代码：

import numpy as np

# 初始化原数组
a = np.array([[[0,1,2],[3,4,5]],
              [[6,7,8],[9,10,11]]])
# 调整b的数值和预期c完全匹配
b = np.array([[2,3],
              [1,1]])

# 生成可广播的z轴索引，形状为(1,1,3)
z_total = a.shape[-1]
z_idx = np.arange(z_total).reshape(1, 1, -1)
# 逐元素生成保留掩码
keep_mask = z_idx < b[..., None]

2. 一次性提取元素并拆分 #

用掩码直接扁平化拉出所有要留的元素，靠累积和算好每个(x,y)位置的切片分割点，NumPy底层直接完成拆分，没有Python层循环：

# 一次性提取所有需要保留的元素
kept = a[keep_mask]
# 计算每个切片的分割位置
split_pos = np.cumsum(b.ravel())
# 按分割点把扁平化的元素拆成对应长度的块
slices = np.split(kept, split_pos)[:-1]
# 把这些块组装成和原数组前两维结构一致的对象数组，存储变长切片
c = np.array(slices, dtype=object).reshape(b.shape)

运行后得到的c和预期结构完全一致：

[[array([0, 1]) array([3, 4, 5])]
 [array([6])    array([9])]]

固定长度数组的常用写法 #

如果不需要存储变长序列，需要结果是标准的三维NumPy数组，可以用掩码将不需要的位置填充为指定值（比如0），生成等长的结果数组：

# 初始化填充为0的同形状数组
c_fixed = np.zeros_like(a)
# 把需要保留的元素按位置填进去
c_fixed[keep_mask] = a[keep_mask]
# 截断到实际需要的最大z长度，去掉冗余的填充维度
c_fixed = c_fixed[..., :b.max()]

性能说明 #

所有核心操作都是NumPy在C层面实现的向量化逻辑，没有Python解释器的逐元素循环开销，处理大数组时速度比写循环快几个数量级。

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