如何用NumPy向量化实现数组元素的查找赋值循环？

你可以利用NumPy的**花式索引（Fancy Indexing）**彻底替换嵌套循环，大幅提升运行效率。具体实现如下：

优化后的完整代码 #

import numpy as np
import cv2

colors = [
    [0, 255, 0],    # GREEN
    [255, 0, 0],    # BLUE
    [0, 0, 255],    # RED
]

WIDTH = 400
HEIGHT = 400
# 将颜色列表转为NumPy数组，适配向量化索引逻辑
colors_np = np.array(colors, dtype=np.uint8)
color_lookup = np.random.randint(0, len(colors), (HEIGHT, WIDTH))

# 一步完成颜色映射，替代双层循环
target = colors_np[color_lookup]

cv2.imshow("target", target)
cv2.waitKey(0)

核心原理 #

• 先把colors列表转为NumPy数组colors_np，让NumPy可以利用底层高效的索引机制处理。
• 直接用二维的color_lookup数组作为索引，NumPy会自动将每个位置的索引值对应到colors_np的行，最终输出形状为(HEIGHT, WIDTH, 3)的数组，完美匹配原target的结构。
• 这种向量化操作在C语言层面执行，完全避免了Python循环的额外开销，数组尺寸越大，性能提升越显著。

可选写法（保留初始化逻辑） #

如果需要保留原代码中先创建空数组再赋值的习惯，也可以这样写：

target = np.zeros((HEIGHT, WIDTH, 3), dtype=np.uint8)
target[:] = colors_np[color_lookup]

效果和直接赋值完全一致，适合需要预先初始化数组的场景。

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