加速NumPy整数数组深度索引：向量化替代循环优化

当然可以！Python层面的循环在处理大数组时速度瓶颈非常明显，用NumPy的向量化广播操作能充分利用底层优化，直接把速度拉满。

先给你直接能用的向量化实现，完美匹配你要的输出格式：

import numpy as np

def im2segmap(im, depth):
    # 生成(depth, H, W)形状的二进制掩码张量
    return (np.arange(depth)[:, np.newaxis, np.newaxis] == im).astype(np.int32)

工作原理拆解 #

咱们一步步看为什么这行代码能替代循环：

1. 维度扩展：np.arange(depth)[:, np.newaxis, np.newaxis]把0到depth-1的一维数组（形状(depth,)）扩展成三维的(depth, 1, 1)结构，这样就能和输入的2D数组im（形状(H, W)）触发NumPy的广播机制。
2. 广播比较：当我们用== im比较时，NumPy会自动把(depth,1,1)和(H,W)广播成(depth, H, W)的数组——每个深度层i都会和整个im数组逐一比较，得到的布尔值就表示im中对应位置是否等于i。
3. 类型转换：最后用astype(np.int32)把布尔值（True/False）转成整数1/0，正好得到你需要的二进制掩码格式。

验证示例输入 #

用你给出的测试输入跑一遍：

im = np.array([[0,2,1],[1,0,1],[2,1,1]])
seg_map = im2segmap(im, 3)
print(seg_map)

输出和你要的目标完全一致：

[[[1 0 0]
  [0 1 0]
  [0 0 0]]

 [[0 0 1]
  [1 0 1]
  [0 1 1]]

 [[0 1 0]
  [0 0 0]
  [1 0 0]]]

性能优势 #

这种实现完全避开了Python循环的开销，所有运算都在NumPy的C底层完成。比如处理1000×1000的大数组，循环版本可能需要几十毫秒，而向量化版本只需要几毫秒，速度提升非常显著。

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